import "ir.hff";
import "common.hff";
import "map.hff";

extern fn ir_genstatic(IR *IRCtx, decl *Decl) void {
   g_targ.B.genstaticf(g_targ.B, decl);
}

struct IntTraits {
   fn hash(a int) u32 { return (a * 7128347) + 972183; }
   fn eq(a int, b int) bool { return a == b; }
}

struct InstStream {
   IR *IRCtx,
   head *IRInst,
   tail *IRInst,
   loopbrks Map<int, *IRInst, IntTraits>,
   loopconts Map<int, *IRInst, IntTraits>,

   fn push(S *InstStream, inst *IRInst) void {
      if S.tail {
         S.tail.next = inst;
      } else {
         S.head = inst;
      }
      S.tail = inst;
   }

   fn mkinst(S *InstStream, inst0 IRInst, narg uint) *IRInst {
      let inst *IRInst = S.IR.alloc->alloc(sizeof IRInst + (narg * sizeof IRArg), alignof IRInst);
      *inst = inst0;
      return inst;
   }

   fn mkinst1(S *InstStream, inst0 IRInst, a IRArg) *IRInst {
      let inst = S->mkinst(inst0, 1);
      inst.args[0] = a;
      return inst;
   }

   fn mkinst2(S *InstStream, inst0 IRInst, a IRArg, b IRArg) *IRInst {
      let inst = S->mkinst(inst0, 2);
      inst.args[0] = a;
      inst.args[1] = b;
      return inst;
   }

   fn mkinst3(S *InstStream, inst0 IRInst, a IRArg, b IRArg, c IRArg) *IRInst {
      let inst = S->mkinst(inst0, 3);
      inst.args[0] = a;
      inst.args[1] = b;
      inst.args[2] = c;
      return inst;
   }

   fn pushnop(S *InstStream) *IRInst  {
      S->push(S->mkinst({:nop}, 0));
      return S.tail;
   }
}


fn genblock(S *InstStream, block [#]Stmt) void;
fn genexpr(S *InstStream, ex *Expr) IRValue {
   static tmpid int = 0;

   fold(ex);

   defmacro genunop(u,t) [ {
      let child = genexpr(S, (u).ex);
      let res = IRValue{ex.ty, .u: :Tmp(tmpid++)};
      let inst = S->mkinst2({t}, :Val(res), :Val(child));
      S->push(inst);
      return res;
   } ]

   defmacro genbinop(b,t) [ {
      let lhs = genexpr(S, (b).lhs);
      let rhs = genexpr(S, (b).rhs);
      let res = IRValue{ex.ty, .u: :Tmp(tmpid++)};
      let inst = S->mkinst3({t}, :Val(res), :Val(lhs), :Val(rhs));
      S->push(inst);
      return res;
   } ]

   switch ex.u {
   case IntLit i; return {ex.ty, :IImm(i.i)};
   case FloLit f; return {ex.ty, :FImm(f)};
   case StrLit s; return {ex.ty, :SImm(s)};
   case BoolLit b; return {ex.ty, :BImm(b)};
   case NullLit; return {ex.ty, :Null};
   case UnOp u;
      switch u.op {
      case :neg;   genunop(u, ex.ty->is(:Int) ? :neg : :fneg);
      case :not;   genunop(u, :not);
      case :compl; genunop(u, :compl);
      case :deref;
         let child = genexpr(S, (u).ex);
         let res = IRValue{ex.ty, .u: :Tmp(tmpid++)};
         let inst = S->mkinst3({:load}, :Val(res), :Val(child), :Val{ty_int,:IImm(0)});
         S->push(inst);
         return res;
      case :preinc, :predec;
         let one = u.op == :preinc ? 1 : -1;
         switch u.ex.u {
         case Symbol decl;
            let ex = genexpr(S, u.ex);
            let inced = IRValue{ex.ty, .u: :Tmp(tmpid++)};
            let rhs = S->mkinst3({:add}, :Val(inced),
                                         :Val(ex), :Val{ty_int, :IImm(one)});
            S->push(rhs);
            let inst = S->mkinst2({:setvar}, :Val(ex), :Val(inced));
            S->push(inst);
            return inced;
         }
      }
   case BinOp b;
      switch b.op {
      case  '+'; genbinop(b, !ex.ty->is(:Flo) ? :add : :fadd);
      case  '-'; genbinop(b, !ex.ty->is(:Flo) ? :sub : :fsub);
      case  '*'; genbinop(b, !ex.ty->is(:Flo) ? :mul : :fmul);
      case  '/'; genbinop(b, !ex.ty->is(:Flo) ? :div : :fdiv);
      case  '%'; genbinop(b, :mod);
      case  '&'; genbinop(b, :band);
      case  '|'; genbinop(b, :bor);
      case  '^'; genbinop(b, :xor);
      case '<<'; genbinop(b, :lsl);
      case '>>'; genbinop(b, ex.ty.u.Int.sgn ? :asr : :lsr);
      case '=='; genbinop(b, :eq);
      case '!='; genbinop(b, :neq);
      case  '<'; genbinop(b, :lt);
      case '<='; genbinop(b, :lteq);
      case  '>';
         let lhs = genexpr(S, (b).lhs);
         let rhs = genexpr(S, (b).rhs);
         let res = IRValue{ex.ty, .u: :Tmp(tmpid++)};
         let inst = S->mkinst3({:lt}, :Val(res), :Val(rhs), :Val(lhs));
         S->push(inst);
         return res;
      case  '>=';
         let lhs = genexpr(S, (b).lhs);
         let rhs = genexpr(S, (b).rhs);
         let res = IRValue{ex.ty, .u: :Tmp(tmpid++)};
         let inst = S->mkinst3({:lteq}, :Val(res), :Val(rhs), :Val(lhs));
         S->push(inst);
         return res;
      case  '=';
         switch b.lhs.u { // var = ex
         case Symbol decl;
            let r = genexpr(S, b.rhs);
            let inst = S->mkinst2({:setvar}, :Val(genexpr(S, b.lhs)), :Val(r));
            S->push(inst);
            return r;
         case UnOp u; // *a = ex
            assert(u.op == :deref, "deref?? unop=");
            let child = genexpr(S, u.ex);
            let r = genexpr(S, b.rhs);
            let inst = S->mkinst3({:store}, :Val(child), :Val{ty_int,:IImm(0)}, :Val(r));
            S->push(inst);
            return r;
         case Index index; // a[b] = ex
            let indexee = genexpr(S, index.lhs);
            let index = genexpr(S, index.rhs);
            let newval = genexpr(S, b.rhs);
            S->push(S->mkinst3({:store}, :Val(indexee), :Val(index), :Val(newval)));
            return newval;
         case else
            assert(#f, "nyi ="); 
         }
      case '+=', '-=', '*=', '/=', '%=', '&=', '|=', '^=', '<<=', '>>=';
         let o IRInstT = b.op == '+='  ? (ex.ty->is(:Flo) ? :fadd : :add)
                       : b.op == '-='  ? (ex.ty->is(:Flo) ? :fsub : :sub)
                       : b.op == '*='  ? (ex.ty->is(:Flo) ? :fmul : :mul)
                       : b.op == '/='  ? (ex.ty->is(:Flo) ? :fdiv : :div)
                       : b.op == '%='  ? :mod
                       : b.op == '&='  ? :band
                       : b.op == '|='  ? :bor
                       : b.op == '<<=' ? :lsl
                       : b.op == '>>=' ? (ex.ty.u.Int.sgn ? :asr : :lsr) : :nop;
         switch b.lhs.u {
         case Symbol decl; // var += ex
            let lhs = genexpr(S, b.lhs);
            let rhs = genexpr(S, b.rhs);
            let new = IRValue{ex.ty, .u: :Tmp(tmpid++)};
            S->push(S->mkinst3({o}, :Val(new), :Val(lhs), :Val(rhs)));
            let inst = S->mkinst2({:setvar}, :Val(lhs), :Val(new));
            S->push(inst);
            return new;
         }

      }
   case Index idx;
      let i = genexpr(S, idx.lhs);
      let off = genexpr(S, idx.rhs);
      let res = IRValue{ex.ty, .u: :Tmp(tmpid++)};
      let inst = S->mkinst3({:load}, :Val(res), :Val(i), :Val(off));
      S->push(inst);
      return res;

   case Cond c;
      let test = genexpr(S, c.test);
      let beqz = S->mkinst1({:beqz}, :Val(test));
      S->push(beqz);
      let t = genexpr(S, c.t);
      let res = IRValue{ex.ty, :Tmp(tmpid++)};
      let inst = S->mkinst2({:copy}, :Val(res), :Val(t));
      S->push(inst);
      let b = S->mkinst({:b}, 0);
      S->push(b);
      beqz.branch = S->pushnop();
      let f = genexpr(S, c.f);
      let inst = S->mkinst2({:copy}, :Val(res), :Val(f));
      S->push(inst);
      b.branch = S->pushnop();
      return res;

   case Call c;
      let n = c.args.#len;
      let inst = S->mkinst({:call, .call_nargs: n}, n + 2);
      inst.args[1] = :Fn(c.lhs.u.Symbol);
      if c.lhs.u.#tag == :Symbol {
         for let i = 0; i < n; ++i {
            inst.args[i + 2] = :Val(genexpr(S, &c.args[i]));
         }
         S->push(inst);
      } else {
      }
      let res = IRValue{ex.ty, :Tmp(tmpid++)};
      inst.args[0] = :Val(res);
      return res;
   case Symbol decl;
      if decl.u.#tag == :Let {
         return {ex.ty, :Local(decl)};
      } else {
         return {ex.ty, :Global(decl)};
      }
   }
   
   assert(#f, "NYI ex");
}

fn genstmt(S *InstStream, stmt *Stmt) void {
   switch stmt.u {
   case Block block;
      genblock(S, block);
   
   case Expr *ex;
      genexpr(S, ex);

   case Decl decl;
      if decl.u.#tag == :Let and !decl.u.Let.ini->empty() {
         let Let = &decl.u.Let;
         let inst = S->mkinst2({:setvar}, :Val{Let.ty, :Local(decl)},
                                          :Val(genexpr(S, &Let.ini.Some)));
         S->push(inst);
      }

   case If If;
      #{
         %t = <test>
         beqz %t, .L0
            <t>
            b .L1
         .L0:
            <f>
         .L1:
      }
      let t = genexpr(S, &If.test);
      let beqz = S->mkinst1({:beqz}, :Val(t));
      S->push(beqz);
      genblock(S, If.t);
      if If.f.#ptr {
         let b = S->mkinst({:b}, 0);
         S->push(b);
         beqz.branch = S->pushnop();
         genblock(S, If.f);
         b.branch = S->pushnop();
      } else {
         beqz.branch = S->pushnop();
      }

   case While loop;
      #{
         .cont:
            %t = <test>
            beqz %t, .brk
            <body>
            b .cont
         .brk:
      }
      let begin = S->pushnop();
      let t = genexpr(S, &loop.test);
      let beqz = S->mkinst1({:beqz}, :Val(t));
      let b = S->mkinst({:b}, 0);
      let end = S->mkinst({:nop}, 0);
      S->push(beqz);
      S.loopconts->put(loop.id, begin);
      S.loopbrks->put(loop.id, end);
      genblock(S, loop.body);
      S->push(b);
      b.branch = begin;
      S->push(end);
      beqz.branch = end;

   case For loop;
      genblock(S, loop.ini);
      let begin = S->pushnop();
      let t = genexpr(S, &loop.test);
      let beqz = S->mkinst1({:beqz}, :Val(t));
      let b = S->mkinst({:b}, 0);
      let next = S->mkinst({:nop}, 0);
      let end = S->mkinst({:nop}, 0);
      S->push(beqz);
      S.loopconts->put(loop.id, next);
      S.loopbrks->put(loop.id, end);
      genblock(S, loop.body);
      S->push(next);
      if !loop.next->empty() {
         genexpr(S, &loop.next.Some);
      }
      S->push(b);
      b.branch = begin;
      S->push(end);
      beqz.branch = end;

   case Break id;
      let jumpto = S.loopbrks->get(id);
      assert(jumpto != #null, "no brk label %d", id);
      S->push(S->mkinst({:b, .branch: *jumpto},0));

   case Continue id;
      let jumpto = S.loopconts->get(id);
      assert(jumpto != #null, "no cont label %d", id);
      S->push(S->mkinst({:b, .branch: *jumpto},0));

   case Return retex;
      if retex->empty() {
         S->push(S->mkinst({:ret0}, 0));
      } else {
         let inst = S->mkinst1({:ret}, :Val(genexpr(S, &retex.Some)));
         S->push(inst);
      }
      
   case else
      assert(#f, "NYI st");
   }
}

fn genblock(S *InstStream, block [#]Stmt) void {
   foreach_ptr(st, _, block) {
      genstmt(S, st);
   }
}

extern fn ir_genfn(IR *IRCtx, f *Fn) void {
   let stream InstStream = {.IR: IR};
   defer stream.loopbrks->clear();
   defer stream.loopconts->clear();

   genblock(&stream, f.body.Some);
   stream->push(stream->mkinst({:ret0}, 0));

   efmt("----------------------\n");
   efmt("function %s.%d:\n", container_of(f, Decl, u.Fn).name, f.id);
   irdump(stream.head);
}