#include "ir.h"

/**  This pass adds in ABI arguments/returns register mappings
 **  and lowers aggregate params/args/returns into scalars
 **
 **  invariant: all `call` instructions when doing this pass shall be preceded by
 **  exactly narg `arg` instructions with no other instructions in between
 **/

struct abiargsvec { vec_of(struct abiarg); };

static int
abiret(struct abiarg abiret[2], struct abiargsvec *abiargs, int *ni, union irtype retty)
{
   short r[2];
   uchar cls[2];
   int retreg = 0;
   if (retty.isagg) {
      retreg = mctarg->abiret(r, cls, ni, retty);
      if (!retreg) {
         vpush(abiargs, ((struct abiarg) { {.cls = KPTR}, r[1] }));
         if (r[0] == -1) {
            memset(abiret, 0, 2*sizeof *abiret);
         } else {
            abiret[0].ty = (union irtype) {.cls = KPTR};
            abiret[0].reg = r[0];
         }
      }
   } else if (retty.cls) {
      retreg = mctarg->abiret(r, cls, ni, retty);
      assert(retreg == 1);
   }
   for (int i = 0; i < retreg; ++i) {
      abiret[i].ty = (union irtype) {.cls = cls[i]};
      abiret[i].reg = r[i];
   }
   return retreg;
}

static int
abiarg(struct abiargsvec *abiargs, int *ni, int *nf, int *ns, union irtype ty)
{
   short r[2];
   uchar cls[2];
   int ret = mctarg->abiarg(r, cls, ni, nf, ns, ty);
   if (!ret) { /* aggregate in stack */
      vpush(abiargs, ((struct abiarg) { ty, -1 }));
   } else if (ret == 1 && ty.isagg && cls[0] == KPTR) { /* aggregate by pointer */
      vpush(abiargs, ((struct abiarg) { {.cls = cls[0]}, r[0] }));
   } else {
      vpush(abiargs, ((struct abiarg) { {.cls = cls[0]}, r[0] }));
      if (ret == 2)
         vpush(abiargs, ((struct abiarg) { {.cls = cls[1]}, r[1] }));
   }
   return ret;
}

/* RPARAM can only appear in the entry block (prologue), each RARG can only appear once.
 * this function patches param starting at instruction no. *start according to cls
 * to patch it to use arg no. `to' (and maybe also `to + 1') */
static void
patchparam(struct function *fn, int *start, int param, int tydat, int to, struct abiarg abi[2])
{
   struct block *blk = fn->entry;
   assert(!blk->phi.n);
   while((*start)++ < blk->ins.n) {
      struct instr *ins = &instrtab[blk->ins.p[*start - 1]];
      if (ins->op == Ocopy && ins->l.t == RPARAM && ins->l.i == param) {
         /* originally aggregate argument */
         assert(tydat != -1);
         if (abi[0].ty.isagg /* aggregate in stack */
          || abi[0].ty.cls == KPTR) /* aggregate by pointer */
         {
            ins->l.i = to;
         } else { /* aggregate in registers */
            const struct typedata *td = &typedata[tydat];
            /* transform
             *  %x = copy %argX
             * into
             *  %x = alloca...
             *  store* %x, %argN
             *  store* %x + I, %argM
             */
            assert(td->siz <= 16 && td->align <= 16);
            ins->op = Oalloca8 + (td->align == 16);
            ins->l = mkref(RICON, td->align == 16 ? 1 : td->siz / 8);
            insertinstr(blk, *start, mkinstr(Ostore1 + ilog2(cls2siz[abi[0].ty.cls]), 0,
                                          mkref(RTMP, ins - instrtab), mkref(RPARAM, to)));
            *start += 1;
            if (td->siz > 8) {
               struct instr tmp = mkinstr(Oadd, KPTR,
                                          mkref(RTMP, ins - instrtab), mkref(RICON, cls2siz[abi[0].ty.cls]));
               insertinstr(blk, *start+1, mkinstr(Ostore1 + ilog2(cls2siz[abi[1].ty.cls]), 0,
                           insertinstr(blk, *start, tmp), mkref(RPARAM, to+1)));
               *start += 2;
            }
         }
         break;
      } else if (oisstore(ins->op) && ins->r.t == RPARAM && ins->r.i == param) {
         /* normal scalar argument */
         assert(tydat == -1);
         ins->r.i = to;
         break;
      }
   }
}

static int
patcharg(struct block *blk, int *icall, struct call *call,
         int argidx, int nabi, struct abiarg abi[2])
{
   int arginst = *icall - (call->narg - argidx);
   struct instr *arg = &instrtab[blk->ins.p[arginst]];
   assert(arg->op == Oarg && arg->l.t == RTYPE);
   if (ref2type(arg->l).isagg) {
      /* aggregate argument */
      if (abi[0].ty.isagg  /* aggregate in stack */
       || abi[0].ty.cls == KPTR) /* aggregate by pointer */
      {
         return 1;
      } else { /* aggregate in registers */
         union ref src = arg->r;
         /* deconstruct into
          * %a = load* %x
          * (%b = load* %x + N)
          */
         delinstr(blk, arginst);
         for (int i = 0; i < nabi; ++i) {
            /* XXX this can generate unaligned loads */
            struct instr ins = {0};
            union ref temp;
            switch (ins.cls = abi[i].ty.cls) {
            default: assert(0);
            case KI4: ins.op = Oloadu4; break;
            case KI8: ins.op = Oloadi8; break;
            case KF4: ins.op = Oloadf4; break;
            case KF8: ins.op = Oloadf8; break;
            }
            if (i == 0)
               ins.l = src;
            else
               ins.l = insertinstr(blk, arginst++,
                                   mkinstr(Oadd, KPTR, src,
                                           mkref(RICON, cls2siz[abi[0].ty.cls])));
            temp = insertinstr(blk, arginst++, ins);
            insertinstr(blk, arginst++, mkarginstr(abi[i].ty, temp));
         }
         *icall += arginst - (call->narg - argidx);
         return nabi;
      }
   } else {
      /* normal scalar argument */
      return 1;
   }
}

void
abi0(struct function *fn)
{
   uint nparam = typedata[fn->fnty.dat].nmemb;
   const union type *paramty = typedata[fn->fnty.dat].param;
   static struct abiarg abiargsbuf[32];
   struct abiargsvec abiargs = {VINIT(abiargsbuf, arraylength(abiargsbuf))};
   int rvovar = -1;
   int ni = 0, nf = 0, ns = 0, istart = 0;
   struct block *blk;
   union ref sret;
   bool sretarghidden = 0;

   if (fn->retty.t == TYVOID) {
      fn->nabiret = 0;
   } else {
      fn->nabiret = abiret(fn->abiret, &abiargs, &ni, mkirtype(fn->retty));
      if (!fn->nabiret && isagg(fn->retty)) { /* ret by hidden pointer */
         sretarghidden = ni == 0;
         sret = insertinstr(fn->entry, 0, mkinstr(Ocopy, KPTR, mkref(RPARAM, 0)));
         ++istart;
      }
   }

   /* adjust params */
   for (int i = 0; i < nparam; ++i) {
      union irtype pty = mkirtype(paramty[i]);
      int thisi = sretarghidden + ni + nf + ns;
      int first = abiargs.n;
      int ret = abiarg(&abiargs, &ni, &nf, &ns, pty);
      if (i != thisi || (pty.isagg && ret))
         patchparam(fn, &istart, i, pty.isagg ? pty.dat : -1, thisi, &abiargs.p[first]);
   }
   fn->abiarg = alloccopy(&fn->arena, abiargs.p, abiargs.n * sizeof *abiargs.p, 0);
   fn->nabiarg = abiargs.n;
   vfree(&abiargs);

   if (!fn->nabiret && isagg(fn->retty)) {
      /* for structures returned by hidden pointer argument,
       * if all return instrs return local var X, make X point to the result location,
       * (return value optimization (RVO)) */
      blk = fn->entry;
      do {
         union ref arg = blk->jmp.arg[0];
         if (blk->jmp.t != Jret) continue;
         if (!arg.t) continue;
         if (arg.t != RTMP || !oisalloca(instrtab[arg.i].op)) {
            rvovar = -1;
            break;
         }
         if (rvovar == -1) {
            rvovar = arg.i;
         } else if (arg.i != rvovar) {
            rvovar = -1;
            break;
         }
      } while ((blk = blk->lnext) != fn->entry);
      if (rvovar != -1)
         instrtab[rvovar] = mkinstr(Ocopy, KPTR, sret);
   }

   blk = fn->entry->lnext;
   do {
      /* adjust calls */
      for (int iinstr = 0; iinstr < blk->ins.n; ++iinstr) {
         struct instr *ins = &instrtab[blk->ins.p[iinstr]];
         struct call *call = &calltab.p[ins->r.i];
         bool structbyval;
         int vararg;

         if (ins->op != Ocall) continue;

         vararg = call->vararg;
         vinit(&abiargs, abiargsbuf, arraylength(abiargsbuf));
         if (!(structbyval = call->ret.isagg))
            for (int i = iinstr - call->narg; i < iinstr; ++i)
               if ((structbyval = ref2type(instrtab[blk->ins.p[i]].l).isagg))
                     break;

         ni = nf = ns = 0;
         memset(call->abiret, 0, sizeof call->abiret);
         if (call->ret.isagg) { /* adjust struct return */
            union ref temp;
            union irtype retty = call->ret;
            struct instr alloca = { .cls = KPTR };
            struct typedata *td = &typedata[retty.dat];
            int ret = abiret(call->abiret, &abiargs, &ni, retty);

            sretarghidden = ni == 0;
            alloca.op = Oalloca8 + (td->align == 16);
            alloca.l = mkref(RICON, td->align == 16 ? 1 : td->siz / 8);
            temp = insertinstr(blk, iinstr++ - call->narg, alloca);
            replref(fn, blk, iinstr, mkref(RTMP, ins - instrtab), temp);
            if (!ret) { /* hidden pointer argument */
               insertinstr(blk, iinstr++ - call->narg,
                           mkarginstr((union irtype){.cls = KPTR}, temp));
               ins->cls = 0;
            } else { /* aggregate returned in regs */
               union ref call2r;
               int to = iinstr + 1;
               assert(in_range(ret, 1, 2));
               ins->cls = call->abiret[0].ty.cls;
               if (ret == 2)
                  call2r = insertinstr(blk, to++, mkinstr(Ocall2r, call->abiret[1].ty.cls,
                                                          mkref(RTMP, ins - instrtab)));
               for (int i = 0; i < ret; ++i) {
                  uchar cls;
                  struct instr store = {0};
                  /* XXX this can generate unaligned stores */
                  switch (cls = call->abiret[i].ty.cls) {
                  default: assert(0);
                  case KF4: case KI4: store.op = Ostore4; break;
                  case KI8: case KF8: store.op = Ostore8; break;
                  }
                  if (i == 0) {
                     store.l = temp;
                     store.r = mkref(RTMP, ins - instrtab);
                  } else {
                     store.l = insertinstr(blk, to++,
                                         mkinstr(Oadd, KPTR, temp,
                                                 mkref(RICON, cls2siz[call->abiret[0].ty.cls])));
                     store.r = call2r;
                  }
                  insertinstr(blk, to++, store);
               }
            }
         } else if (ins->cls) {
            int ret = abiret(call->abiret, &abiargs, &ni, (union irtype){.cls = ins->cls});
            assert(ret == 1 && !ni);
         }

         for (int i = 0; i < call->narg; ++i) {
            union irtype pty = ref2type(instrtab[blk->ins.p[iinstr - call->narg + i]].l);
            int thisi = sretarghidden + ni + nf + ns;
            int first = abiargs.n;
            int ret = abiarg(&abiargs, &ni, &nf, &ns, pty);
            ret = patcharg(blk, &iinstr, call, i, ret, &abiargs.p[first]);
            if (call->vararg == i)
               vararg = thisi;
         }
         if (call->ret.isagg) call->ret = (union irtype){0};
         call->vararg = vararg;
         call->abiargregs = alloc(&fn->arena, abiargs.n * sizeof *call->abiargregs, 0);
         for (int i = 0; i < abiargs.n; ++i) call->abiargregs[i] = abiargs.p[i].reg;
         call->narg = abiargs.n;
         vfree(&abiargs);
      }

      /* adjust returns */
      if (isagg(fn->retty) && blk->jmp.t == Jret && blk->jmp.arg[0].t) {
         assert(!blk->jmp.arg[1].t);
         if (fn->nabiret) { /* aggregate return in register(s) */
            union ref src = blk->jmp.arg[0];
            for (int i = 0; i < fn->nabiret; ++i) {
               /* XXX this can generate unaligned loads */
               struct instr ins = {0};
               switch (ins.cls = fn->abiret[i].ty.cls) {
               default: assert(0);
               case KI4: ins.op = Oloadu4; break;
               case KI8: ins.op = Oloadi8; break;
               case KF4: ins.op = Oloadf4; break;
               case KF8: ins.op = Oloadf8; break;
               }
               if (i == 0)
                  ins.l = src;
               else
                  ins.l = insertinstr(blk, blk->ins.n,
                                      mkinstr(Oadd, KPTR, src,
                                              mkref(RICON, cls2siz[fn->abiret[0].ty.cls])));
               blk->jmp.arg[i] = insertinstr(blk, blk->ins.n, ins);
            }
         } else {
            /* aggregate return (arg[0] is pointer to return value) */
            if (rvovar == -1) {
               /* blit %sret, %arg */
               union irtype typ = mkirtype(fn->retty);
               insertinstr(blk, blk->ins.n, mkarginstr(typ, sret));
               insertinstr(blk, blk->ins.n, mkarginstr(typ, blk->jmp.arg[0]));
               insertinstr(blk, blk->ins.n, mkintrin(fn, INstructcopy, 0, 2));
            } else assert(blk->jmp.arg[0].bits == mkref(RTMP, rvovar).bits);
            if (fn->abiret[0].ty.cls) blk->jmp.arg[0] = rvovar == -1 ? sret : mkref(RTMP, rvovar);
            else memset(blk->jmp.arg, 0, sizeof blk->jmp.arg);
         }
      }
   } while ((blk = blk->lnext) != fn->entry);

   if (ccopt.dbg.a) {
      efmt("after abi0:\n");
      irdump(fn, fn->name);
   }
}

/* vim:set ts=3 sw=3 expandtab: */