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

extern fn ir_genstatic(IR *IRCtx, decl *Decl) void {
}

struct InstStream;


struct InstStream {
   IR *IRCtx,
   head *IRInst,
   tail *IRInst,
   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) void {
      S->push(S->mkinst({:nop}, 0));
   }
}


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 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 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);
      S->pushnop();
      beqz.branch = S.tail;
      let f = genexpr(S, c.f);
      let inst = S->mkinst2({:copy}, :Val(res), :Val(f));
      S->push(inst);
      S->pushnop();
      b.branch = S.tail;
      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 If If;
      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);
         S->pushnop();
         beqz.branch = S.tail;
         genblock(S, If.f);
         S->pushnop();
         b.branch = S.tail;
      } else {
         S->pushnop();
         beqz.branch = S.tail;
      }

   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};

   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);
}