/* -*- mesa-c++  -*-
 *
 * Copyright (c) 2022 Collabora LTD
 *
 * Author: Gert Wollny <gert.wollny@collabora.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * on the rights to use, copy, modify, merge, publish, distribute, sub
 * license, and/or sell copies of the Software, and to permit persons to whom
 * the Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice (including the next
 * paragraph) shall be included in all copies or substantial portions of the
 * Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL
 * THE AUTHOR(S) AND/OR THEIR SUPPLIERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
 * USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#ifndef SFN_SHADER_H
#define SFN_SHADER_H

#include "sfn_instr.h"
#include "sfn_instrfactory.h"
#include "sfn_instr_controlflow.h"
#include "gallium/drivers/r600/r600_shader.h"
#include "sfn_liverangeevaluator.h"

#include <bitset>
#include <memory>
#include <stack>
#include <vector>

struct nir_shader;
struct nir_cf_node;
struct nir_if;
struct nir_block;
struct nir_instr;

namespace r600 {

class ShaderIO {
public:
   void set_sid(int sid);
   void override_spi_sid(int spi_sid);
   void print(std::ostream& os) const;

   int spi_sid() const { return m_spi_sid;}
   unsigned sid() const { return m_sid;}

   int location() const {return m_location;}
   unsigned name() const { return m_name;}

   int pos() const { return m_pos;}
   void set_pos(int pos) {m_pos = pos;}

   bool is_param() const { return m_is_param;}
   void set_is_param(bool val) { m_is_param = val;}

   void set_gpr(int gpr) {m_gpr = gpr;}
   int gpr() const {return m_gpr;}

protected:
   ShaderIO(const char *type, int loc, int name);

private:

   virtual void do_print(std::ostream& os) const = 0;

   const char *m_type;
   int m_location{-1};
   int m_name{-1};
   int m_sid{0};
   int m_spi_sid{0};
   int m_pos{0};
   int m_is_param{false};
   int m_gpr{0};
};

class ShaderOutput : public ShaderIO {
public:
   ShaderOutput();
   ShaderOutput(int location, int name, int writemask);

   int writemask() const { return m_writemask;}

private:
   void do_print(std::ostream& os) const override;

   int m_writemask{0};
};


class ShaderInput : public ShaderIO  {
public:
   ShaderInput();
   ShaderInput(int location, int name);
   void set_interpolator(int interp, int interp_loc, bool uses_interpolate_at_centroid);
   void set_uses_interpolate_at_centroid();
   void set_need_lds_pos() { m_need_lds_pos = true;}
   int ij_index() const { return m_ij_index;}

   int interpolator() const{return m_interpolator;}
   int interpolate_loc() const {return m_interpolate_loc;}
   bool need_lds_pos() const {return m_need_lds_pos;}
   int lds_pos() const {return m_lds_pos;}
   void set_lds_pos(int pos) {m_lds_pos = pos;}

   int ring_offset() const {return m_ring_offset;}
   void set_ring_offset(int offs) {m_ring_offset = offs;}
   bool uses_interpolate_at_centroid() const {return m_uses_interpolate_at_centroid;}

private:
   void do_print(std::ostream& os) const override;

   int m_interpolator{0};
   int m_interpolate_loc{0};
   int m_ij_index{0};
   bool m_uses_interpolate_at_centroid{false};
   bool m_need_lds_pos{false};
   int m_lds_pos{0};
   int m_ring_offset{0};
};

class Shader : public Allocate {
public:
   using InputIterator = std::map<int, ShaderInput>::iterator;
   using OutputIterator = std::map<int, ShaderOutput>::iterator;

   using ShaderBlocks = std::list<Block::Pointer, Allocator<Block::Pointer>>;

   Shader(const Shader& orig) = delete;

   virtual ~Shader() {}

   bool add_info_from_string(std::istream& is);

   static Shader *translate_from_nir(nir_shader *nir, const pipe_stream_output_info *so_info, r600_shader *gs_shader,
                                     r600_shader_key& key, r600_chip_class chip_class);

   bool process(nir_shader *nir);

   bool process_cf_node(nir_cf_node *node);
   bool process_if(nir_if *node);
   bool process_loop(nir_loop *node);
   bool process_block(nir_block *node);
   bool process_instr(nir_instr *instr);
   void emit_instruction(PInst instr);
   bool emit_atomic_local_shared(nir_intrinsic_instr* instr);

   void print(std::ostream& os ) const;
   void print_header(std::ostream& os ) const;

   bool process_intrinsic(nir_intrinsic_instr *intr);

   virtual bool load_input(nir_intrinsic_instr *intr) = 0;
   virtual bool store_output(nir_intrinsic_instr *intr) = 0;

   bool load_uniform(nir_intrinsic_instr *intr);
   bool load_ubo(nir_intrinsic_instr *intr);

   ValueFactory& value_factory();

   void add_output(const ShaderOutput& output) {
      m_outputs[output.location()] = output;
   }

   void add_input(const ShaderInput& input) {
      m_inputs[input.location()] = input;
   }

   void set_input_gpr(int driver_lcation, int gpr);

   InputIterator find_input(int location) { return m_inputs.find(location);}

   InputIterator input_not_found() {return m_inputs.end();}

   OutputIterator find_output(int location);
   OutputIterator output_not_found() {return m_outputs.end();}

   ShaderBlocks& func() { return m_root; }
   void reset_function(ShaderBlocks& new_root);

   void emit_instruction_from_string(const std::string &s);

   void set_info(nir_shader *nir);
   void get_shader_info(r600_shader *sh_info);

   r600_chip_class chip_class() const {return m_chip_class;};
   void set_chip_class(r600_chip_class cls) {m_chip_class = cls;};

   void start_new_block(int nesting_depth);

   const ShaderOutput& output(int base) const;

   LiveRangeMap prepare_live_range_map();

   void set_last_txd(Instr *txd){m_last_txd = txd;}
   Instr *last_txd(){return m_last_txd;}

   // Needed for keeping the memory access in order
   void chain_scratch_read(Instr *instr);
   void chain_ssbo_read(Instr *instr);

   virtual uint32_t enabled_stream_buffers_mask() const {return 0;};

   size_t noutputs() const { return m_outputs.size();}
   size_t ninputs() const { return m_inputs.size();}

   enum Flags {
      sh_indirect_const_file,
      sh_needs_scratch_space,
      sh_needs_sbo_ret_address,
      sh_uses_atomics,
      sh_uses_images,
      sh_uses_tex_buffer,
      sh_writes_memory,
      sh_txs_cube_array_comp,
      sh_indirect_atomic,
      sh_mem_barrier,
      sh_legacy_math_rules,
      sh_flags_count
   };

   void set_flag(Flags f) {m_flags.set(f);}
   bool has_flag(Flags f) const {return m_flags.test(f);}

   int atomic_file_count() const { return m_atomic_file_count; }

   PRegister atomic_update();
   int remap_atomic_base(int base);
   auto evaluate_resource_offset(nir_intrinsic_instr *instr, int src_id) -> std::pair<int, PRegister>;
   int ssbo_image_offset() const {return m_ssbo_image_offset;}
   PRegister rat_return_address() {assert(m_rat_return_address); return m_rat_return_address;}

   PRegister emit_load_to_register(PVirtualValue src);

protected:
   enum ESlots {
      es_face,
      es_instanceid,
      es_invocation_id,
      es_patch_id,
      es_pos,
      es_rel_patch_id,
      es_sample_mask_in,
      es_sample_id,
      es_sample_pos,
      es_tess_factor_base,
      es_vertexid,
      es_tess_coord,
      es_primitive_id,
      es_helper_invocation,
      es_last
   };

   std::bitset<es_last> m_sv_values;

   Shader(const char *type_id);

   const ShaderInput& input(int base) const;

   bool emit_simple_mov(nir_dest& dest, int chan, PVirtualValue src, Pin pin = pin_free);

   template <typename T>
   using IOMap = std::map<int, T, std::less<int>, Allocator<std::pair<const int, T>>>;

   IOMap<ShaderInput>& inputs() {return m_inputs;}

private:
   virtual bool process_stage_intrinsic(nir_intrinsic_instr *intr) = 0;

   bool allocate_registers_from_string(std::istream& is, Pin pin);
   bool allocate_arrays_from_string(std::istream& is);

   bool read_chipclass(std::istream& is);

   bool load_uniform_indirect(nir_intrinsic_instr *intr, PVirtualValue addr, int offset , int buffer_id);

   bool scan_shader(const nir_function *impl);
   bool scan_uniforms(nir_variable *uniform);
   void allocate_reserved_registers();

   void allocate_local_registers(const exec_list *registers);

   virtual int do_allocate_reserved_registers() = 0;

   bool scan_instruction(nir_instr *instr);
   virtual bool do_scan_instruction(nir_instr *instr) = 0;

   void print_properties(std::ostream& os) const;
   virtual void do_print_properties(std::ostream& os) const = 0;

   bool read_output(std::istream& is);
   bool read_input(std::istream& is);
   virtual bool read_prop(std::istream& is) = 0;

   bool emit_if_start(nir_if *if_stmt);
   bool emit_control_flow(ControlFlowInstr::CFType type);
   bool emit_store_scratch(nir_intrinsic_instr *intr);
   bool emit_load_scratch(nir_intrinsic_instr *intr);
   bool emit_local_store(nir_intrinsic_instr *intr);
   bool emit_local_load(nir_intrinsic_instr* instr);
   bool emit_load_tcs_param_base(nir_intrinsic_instr* instr, int offset);
   bool emit_barrier(nir_intrinsic_instr* intr);
   bool emit_shader_clock(nir_intrinsic_instr* instr);
   bool emit_wait_ack();

   bool equal_to(const Shader& other) const;
   void finalize();
   virtual void do_finalize();

   virtual void do_get_shader_info(r600_shader *sh_info);

   ShaderBlocks m_root;
   Block::Pointer m_current_block;

   InstrFactory *m_instr_factory;
   const char *m_type_id;

   IOMap<ShaderOutput> m_outputs;
   IOMap<ShaderInput> m_inputs;
   r600_chip_class m_chip_class;

   int m_scratch_size;
   int m_next_block;
   bool m_indirect_const_file{false};

   Instr *m_last_txd {nullptr};

   uint32_t m_indirect_files{0};      
   std::bitset<sh_flags_count> m_flags;
   uint32_t nhwatomic_ranges{0};
   std::vector<r600_shader_atomic> m_atomics;

   uint32_t m_nhwatomic{0};
   uint32_t m_atomic_base{0};
   uint32_t m_next_hwatomic_loc{0};
   std::unordered_map<int, int> m_atomic_base_map;
   uint32_t m_atomic_file_count{0};
   PRegister m_atomic_update{nullptr};
   PRegister m_rat_return_address{nullptr};

   int32_t m_ssbo_image_offset{0};
   uint32_t m_nloops{0};

   class InstructionChain : public InstrVisitor {
   public:
      void visit(AluInstr  *instr) override {(void) instr;}
      void visit(AluGroup *instr) override {(void) instr;}
      void visit(TexInstr *instr) override {(void) instr;}
      void visit(ExportInstr *instr) override {(void) instr;}
      void visit(FetchInstr *instr) override {(void) instr;}
      void visit(Block *instr) override {(void) instr;}
      void visit(ControlFlowInstr *instr) override {(void) instr;}
      void visit(IfInstr *instr) override {(void) instr;}
      void visit(StreamOutInstr *instr) override {(void) instr;}
      void visit(MemRingOutInstr *instr) override {(void) instr;}
      void visit(EmitVertexInstr *instr) override {(void) instr;}
      void visit(WriteTFInstr *instr) override {(void) instr;}
      void visit(LDSAtomicInstr *instr) override {(void) instr;}
      void visit(LDSReadInstr *instr) override {(void) instr;}

      void visit(ScratchIOInstr *instr) override;
      void visit(GDSInstr *instr) override;
      void visit(RatInstr *instr) override;

      void apply(Instr *current, Instr **last);

      Shader *this_shader{nullptr};
      Instr *last_scratch_instr{nullptr};
      Instr *last_gds_instr{nullptr};
      Instr *last_ssbo_instr{nullptr};
      bool prepare_mem_barrier{false};
   };

   InstructionChain m_chain_instr;
   std::vector<Instr *> m_loops;
};


std::pair<unsigned, unsigned>
r600_get_varying_semantic(unsigned varying_location);

}

#endif // SHADER_H