4 files changed, 83 insertions(+), 13 deletions(-)

M src/Compiler/CPS.idr
M src/Compiler/CPSLift.idr
M src/Compiler/LLVM/Rapid/CPSCodegen.idr
M src/Compiler/PrepareCode.idr

@@ 22,22 22,27 @@ data ContName : Type where
   BigK : Int -> ContName
   -- local continuation
   MkContName : String -> Int -> ContName
+  -- closure as continuation
+  Clos : Name -> ContName
 
+export
 Show ContName where
   show (BigK i) = "[K\{show i}]"
   show (MkContName s i) = "K:" ++ s ++ ":" ++ (show i)
+  show (Clos n) = "K:<\{show n}>"
 
 Eq ContName where
   (BigK i) == (BigK j) = (i == j)
   (MkContName s1 i1) == (MkContName s2 i2) =
     (i1 == i2) && (s1 == s2)
+  (Clos n1) == (Clos n2) = n1 == n2
   _ == _ = False
 
 mutual
   public export
   data CPSNmExp : Type where
     KNAppCont : FC -> (k : ContName) -> (arg : Atom) -> CPSNmExp
-    KNJump : FC -> (f : Atom) -> (kn : ContName) -> (args : List Atom) -> CPSNmExp
+    KNJump : FC -> (f : Atom) -> (k : ContName) -> (args : List Atom) -> CPSNmExp
     KNCon : FC -> Name -> ConInfo -> (tag : Maybe Int) -> List Atom -> Name -> CPSNmExp -> CPSNmExp
     KNExtPrim : FC -> Name -> List Atom -> Name -> CPSNmExp -> CPSNmExp
     KNFix : FC -> (k : ContName) -> (args : List Name) -> (body : CPSNmExp) -> Name -> CPSNmExp -> CPSNmExp


@@ 168,9 173,12 @@ cps nm (NmApp fc f args) kf = do
   resultVar <- getUnique' "res_"
   cont <- kf (KLocal emptyFC resultVar)
   contName <- makeContName "app"
+  bk <- makeBigK
+  contName' <- getUnique' "app"
   app <- cps nm f (\x =>
-    argHelper nm args (\cargs => pure $ KNJump fc x contName cargs))
-  pure (KNLetCont fc resultVar cont contName app)
+    argHelper nm args (\cargs => pure $ KNJump fc x (Clos contName') cargs))
+  --pure (KNLetCont fc resultVar cont contName app)
+  pure (KNFix fc bk [resultVar] cont contName' app)
 
 cps nm (NmCon fc n ci tag args) kf = do
   resultVar <- getUnique' "con_"


@@ 193,7 201,7 @@ cps nm (NmConstCase fc sc alts def) kf = do
   lambdaArg <- getUnique' "cca_"
   cbody <- kf (KLocal emptyFC lambdaArg)
 
-  let contSmall = isSmallerThan 2 cbody
+  let contSmall = False --isSmallerThan 2 cbody
   let usekf = if contSmall then kf else newkf
 
   cps nm sc (\csc => do


@@ 208,7 216,7 @@ cps nm (NmConCase fc sc alts def) kf = do
   lambdaArg <- getUnique' "cca_"
   cbody <- kf (KLocal emptyFC lambdaArg)
 
-  let contSmall = isSmallerThan 5 cbody
+  let contSmall = False --isSmallerThan 5 cbody
   let usekf = if contSmall then kf else newkf
 
   cps nm sc (\csc => do


@@ 224,11 232,13 @@ cps nm (NmExtPrim fc p args) kf = do
     pure $ KNExtPrim fc p cargs resultVar cont)
 -- TODO: test:
 cps nm (NmForce fc reason e) kf = do
-  contName <- makeContName "tk"
+  contName <- getUnique' "tk"
+  bk <- makeBigK
   resultVar <- getUnique' "fres_"
   cont <- kf (KLocal emptyFC resultVar)
   cps nm e (\thunk => do
-    pure $ KNLetCont fc resultVar cont contName $ KNJump fc thunk contName [])
+    --pure $ KNLetCont fc resultVar cont contName $ KNJump fc thunk contName [])
+    pure $ KNFix fc bk [resultVar] cont contName $ KNJump fc thunk (Clos contName) [])
 -- TODO: test:
 cps nm (NmDelay fc reason e) kf = do
   thunkCont <- makeBigK


@@ 341,7 351,7 @@ khalt arg = pure $ KNAppCont emptyFC (BigK 0) arg
 
 toCPS : NamedCExp -> Core CPSNmExp
 toCPS e = do
-  cpsIdx <- newRef CPSIndex 0
+  cpsIdx <- newRef CPSIndex 1
   cps empty e khalt
 
 simplifyTwice : CPSNmExp -> CPSNmExp

@@ 19,6 19,8 @@ mutual
     KLJump : FC -> (f : Name) -> (args : List Atom) -> CPSLifted
     KLApp : FC -> (f : Atom) -> (args : List Atom) -> CPSLifted
     KLClosure : FC -> Name -> (missing : Nat) -> (free : List Atom) -> Name -> CPSLifted -> CPSLifted
+    KLLocalCont : FC -> (arg : Name) -> (body : CPSLifted) -> ContName -> CPSLifted -> CPSLifted
+    KLAppCont : FC -> ContName -> (arg : Atom) -> CPSLifted
     KLCon : FC -> Name -> ConInfo -> (tag : Maybe Int) -> List Atom -> Name -> CPSLifted -> CPSLifted
     KLExtPrim : FC -> Name -> List Atom -> Name -> CPSLifted -> CPSLifted
     KLOp : {arity : _} -> FC -> PrimFn arity -> Vect arity Atom -> Name -> CPSLifted -> CPSLifted


@@ 49,6 51,8 @@ mutual
   Show CPSLifted where
     show (KLJump _ n args) = assert_total "(KLJump \{show n} [\{show args}])"
     show (KLApp _ f args) = assert_total "(KLApp \{show f} [\{show args}])"
+    show (KLLocalCont _ arg body kn k) = assert_total "(KLLocalCont [\{show arg}] => \{show body}, bind \{show kn} in \{show k})"
+    show (KLAppCont _ kn arg) = assert_total "(KLAppCont \{show kn} [\{show arg}])"
     show (KLClosure _ fn miss args n k) = assert_total "(Closure \{show fn} (\{show args}), bind \{show n} in \{show k} )"
     show (KLCon _ name ci tag args n k) = assert_total "[%con \{show ci} \{show name} \{show tag} \{show args}, bind \{show n} in \{show k}]"
     show (KLExtPrim _ prim args n k) = assert_total "[%extern \{show prim} \{show args}, bind \{show n} in \{show k}]"


@@ 73,6 77,8 @@ mutual
   ishow : CPSLifted -> String
   ishow (KLJump _ n args) = assert_total "(KLJump \{show n} [\{show args}])"
   ishow (KLApp _ f args) = assert_total "(KLApp \{show f} [\{show args}])"
+  ishow (KLLocalCont _ arg body kn k) = assert_total "(KLLocalCont (\{show arg}) => (\{ishow body}), bind \{show kn} in ... )"
+  ishow (KLAppCont _ kn arg) = assert_total "(KLAppCont \{show kn} [\{show arg}])"
   ishow (KLClosure _ fn miss args n k) = assert_total "(Closure \{show fn} (\{show args}), bind \{show n} in ... )"
   ishow (KLCon _ name ci tag args n k) = assert_total "[%con \{show ci} \{show name} \{show tag} \{show args}, bind \{show n} in ... ]"
   ishow (KLExtPrim _ prim args n k) = assert_total "[%extern \{show prim} \{show args}, bind \{show n} in ... ]"


@@ 85,6 91,12 @@ mutual
 SortedSet : Type -> Type
 SortedSet k = SortedMap k ()
 
+export
+decontName : ContName -> Name
+decontName (BigK i) = MN ("__cps_BIGK") i
+decontName (MkContName s i) = MN ("__cps_" ++ s) i
+decontName (Clos n) = n
+
 findFreeAtom : Atom -> SortedSet Name
 findFreeAtom (KLocal _ n) = singleton n ()
 findFreeAtom _ = empty


@@ 98,6 110,15 @@ findFree (KLJump _ f args) =
   foldMap findFreeAtom args
 findFree (KLApp _ f args) =
   findFreeAtom f <+> foldMap findFreeAtom args
+findFree (KLAppCont _ kn@(MkContName _ _) arg) =
+  -- ignore local continuation
+  findFreeAtom arg
+findFree (KLAppCont _ kn arg) =
+  insert (decontName kn) () $ findFreeAtom arg
+findFree (KLLocalCont _ arg body _ k) =
+  let freeK = findFree k in
+  let freeBody = delete arg (findFree body) in
+  freeK <+> freeBody
 findFree (KLClosure _ _ _ args n k) =
   let freeK = delete n (findFree k) in
   freeK <+> foldMap findFreeAtom args


@@ 145,10 166,6 @@ emitFunc fc free args body = do
   put Lifted ((n+1), f::funcs)
   pure newName
 
-decontName : ContName -> Name
-decontName (BigK i) = MN ("__cps_BIGK") i
-decontName (MkContName s i) = MN ("__cps_" ++ s) i
-
 liftCPS : {auto lif : Ref Lifted LiftedState} ->
           CPSNmExp ->
           Core CPSLifted


@@ 163,9 180,12 @@ liftCPS (KNLetCont fc arg body kn k) = do
   lbody <- liftCPS body
   lk <- liftCPS k
   let rawKn = decontName kn
+  {-
   let freeVars = keys $ delete rawKn $ delete arg (findFree lbody)
   liftedFuncName <- emitFunc fc freeVars [arg] lbody
   pure (KLClosure fc liftedFuncName 1 (map (KLocal emptyFC) freeVars) rawKn lk)
+  -}
+  pure (KLLocalCont fc arg lbody kn lk)
 liftCPS (KNCon fc name ci tag args n k) = do
   lk <- liftCPS k
   pure (KLCon fc name ci tag args n lk)


@@ 202,7 222,7 @@ liftCPS (KNJump fc (KLocal fc' n) kn args) = do
 liftCPS (KNJump fc x _ _) = do
   pure (KLCrash fc "cannot call \{show x}")
 liftCPS (KNAppCont fc kn arg) = do
-  pure (KLApp fc (KLocal emptyFC (decontName kn)) [arg])
+  pure (KLAppCont fc kn arg)
 liftCPS (KNLetConst fc c n k) = do
   lk <- liftCPS k
   pure (KLLetConst fc c n lk)

@@ 692,6 692,7 @@ getInstIR instr@(KLClosure fc fn missing free n k) = do
   appendCode "; \{ishow instr}"
   mkClosure fn n (cast missing) free
   getInstIR k
+
 getInstIR instr@(KLJump fc fn args) =
   if length args <= MAX_ARGS_DIRECT
   then do


@@ 709,6 710,40 @@ getInstIR instr@(KLJump fc fn args) =
     let iicName = MN "__iic" !(getUnique)
     appendCode "%\{safeName dummyArgName} = bitcast i8* null to %ObjPtr\n"
     getInstIR (KLClosure fc fn 1 args iicName (KLApp fc (KLocal fc iicName) [KLocal fc dummyArgName]))
+getInstIR instr@(KLLocalCont fc arg body kn k) = do
+  let lbl = "k" ++ (safeName $ decontName kn)
+  appendCode "; prepare arg for \{lbl}:"
+  appendCode "%\{lbl}.arg0 = alloca ptr"
+  getInstIR k
+  appendCode "; \{ishow instr}"
+  appendCode "\{lbl}:"
+  appendCode "; load arg for \{lbl}:"
+  appendCode "%\{safeName arg} = load ptr, ptr %\{lbl}.arg0"
+  getInstIR body
+getInstIR instr@(KLAppCont fc (BigK n) arg) = do
+  appendCode "; \{ishow instr}"
+  heapOffset <- getHeapOffset
+  when (heapOffset > 0) (addError "more heap allocated than checked")
+  let heapCheck = SSA I64 "%HeapCheck"
+  heapNext <- mkSub heapCheck (pConst (-heapOffset))
+  appendCode $ "musttail call tailcc void @__apply_closure1(" ++ (toIR heapNext) ++ ", %TSOPtr %BaseArg, ptr %\{safeName $ decontName (BigK n)}, " ++ (toIR $ !(atom2val arg)) ++ ")"
+  appendCode "ret void"
+getInstIR instr@(KLAppCont fc (Clos n) arg) = do
+  appendCode "; \{ishow instr}"
+  heapOffset <- getHeapOffset
+  when (heapOffset > 0) (addError "more heap allocated than checked")
+  let heapCheck = SSA I64 "%HeapCheck"
+  heapNext <- mkSub heapCheck (pConst (-heapOffset))
+  appendCode $ "musttail call tailcc void @__apply_closure1(" ++ (toIR heapNext) ++ ", %TSOPtr %BaseArg, ptr %\{safeName n}, " ++ (toIR $ !(atom2val arg)) ++ ")"
+  appendCode "ret void"
+getInstIR instr@(KLAppCont fc kn arg) = do
+  appendCode "; \{ishow instr}"
+  heapOffset <- getHeapOffset
+  when (heapOffset > 0) (addError "more heap allocated than checked")
+  let heapCheck = SSA I64 "%HeapCheck"
+  heapNext <- mkSub heapCheck (pConst (-heapOffset))
+  appendCode $ "store ptr %k\{safeName $ decontName kn}.arg0, " ++ toIR !(atom2val arg)
+  appendCode $ "br label %k\{safeName $ decontName kn}"
 getInstIR instr@(KLApp fc f [arg1]) = do
   appendCode "; \{ishow instr}"
   heapOffset <- getHeapOffset


@@ 1002,6 1037,9 @@ maxRequiredHeapSpace (KLJump _ _ args) = let argCount = length args in
   if argCount <= MAX_ARGS_DIRECT then 0
     else 16 + 8 * (cast {to=Int} argCount)
 maxRequiredHeapSpace (KLApp _ _ _) = 0
+maxRequiredHeapSpace (KLAppCont _ _ _) = 0
+maxRequiredHeapSpace (KLLocalCont _ _ body _ k) =
+  maxRequiredHeapSpace body + maxRequiredHeapSpace k
 maxRequiredHeapSpace (KLClosure _ _ _ args _ k) = 16 + 8 * (cast {to=Int} $ length args) + maxRequiredHeapSpace k
 maxRequiredHeapSpace (KLCon _ _ _ _ args _ k) = 8 + 8 * (cast {to=Int} $ length args) + maxRequiredHeapSpace k
 maxRequiredHeapSpace (KLExtPrim _ (NS _ primName) _ _ k) =

@@ 16,7 16,9 @@ constructorNamesConstAlt : CLConstAlt -> List Name
 collectConstructorNames : CPSLifted -> List Name
 collectConstructorNames (KLJump _ _ _) = []
 collectConstructorNames (KLApp _ _ _) = []
+collectConstructorNames (KLAppCont _ _ _) = []
 collectConstructorNames (KLClosure _ _ _ _ _ k) = collectConstructorNames k
+collectConstructorNames (KLLocalCont _ _ body _ k) = collectConstructorNames body ++ collectConstructorNames k
 collectConstructorNames (KLCon _ name _ (Just tag) _ _ k) = collectConstructorNames k
 collectConstructorNames (KLCon _ name _ Nothing _ _ k) = name :: collectConstructorNames k
 collectConstructorNames (KLExtPrim _ _ _ _ k) = collectConstructorNames k