Fix TCO and eraseLambdas transforms

* Fix TCO mutation issues by copying to and from the tco vars.  This
  is a bit more complex than the JS equivalent because we cannot
  directly rely on scoped variable re-assignments from some inner
  lambda.
* Also fix the eraseLambdas transform to always generate functions of
  the right shape, namely: (ctx, arg, ...)

runtime/purescript.c

@@ -328,6 +328,12 @@ const ANY * purs_any_concat(const ANY * x, const ANY * y) {
 	}
 }
 
+inline const ANY * purs_any_copy(const ANY * src) {
+	ANY * copy = purs_new(const ANY);
+	memcpy(copy, src, sizeof (const ANY));
+	return copy;
+}
+
 // -----------------------------------------------------------------------------
 // strings
 // -----------------------------------------------------------------------------

runtime/purescript.h

@@ -156,6 +156,7 @@ int purs_any_eq_num    (const ANY *, double);
 
 int purs_any_eq(const ANY *, const ANY *);
 const ANY * purs_any_concat(const ANY *, const ANY *);
+const ANY * purs_any_copy(const ANY *);
 
 // -----------------------------------------------------------------------------
 // strings
@@ -256,7 +257,11 @@ const ANY * purs_thunked_deref(const void * data);
 
 #define purs_any_int_neg(X) purs_any_int_new(-purs_any_get_int(X))
 #define purs_any_int_set_mut(X, V) do { X->value.i = V; } while (0)
-#define purs_any_assign_mut(V1, V2) do { ((ANY*) V1)->tag = V2->tag; ((ANY*)V1)->value = V2->value; } while (0)
+#define purs_any_assign_mut(V1, V2)\
+	do {\
+		((ANY*) V1)->tag = V2->tag;\
+		((ANY*) V1)->value = V2->value;\
+	} while (0)
 
 /* code-gen helper to allocate and fill a scope.
  * assumes scope to consist only of (const ANY *) pointers, the count of which

src/Language/PureScript/CodeGen/C/Optimizer/Inliner.purs

@@ -76,14 +76,14 @@ unThunk = AST.everywhere go
   go block@(AST.Block asts) =
     case A.unsnoc asts of
       Just
-        { last:
+        { init
+        , last:
              AST.Return
               (AST.App (AST.Var "purs_any_app")
                 [ AST.Function
                     { arguments: [], body: Just (AST.Block body) }
                 , AST.Null
                 ])
-        , init
         } ->
         AST.Block $ init <> body
       _ ->

src/Language/PureScript/CodeGen/C/Optimizer/TCO.purs

@@ -8,13 +8,12 @@ import Data.Array ((:))
 import Data.Array as A
 import Data.Either (Either(..))
 import Data.Foldable (all)
-import Data.Maybe (Maybe(..), fromMaybe)
+import Data.Maybe (Maybe(..))
 import Data.Newtype (wrap)
 import Data.String as Str
 import Data.Tuple.Nested ((/\))
 import Language.PureScript.CodeGen.C.AST (AST)
 import Language.PureScript.CodeGen.C.AST as AST
-import Language.PureScript.CodeGen.C.AST as Type
 import Language.PureScript.CodeGen.Runtime as R
 
 -- | Eliminate tail calls
@@ -46,10 +45,13 @@ tco = AST.everywhere convert
       else x
   convert x = x
 
-  tcoDone = "$tco_done"
+  tcoDone = "$_tco_done"
   tcoLoop = "$_tco_loop"
   tcoResult = "$_tco_result"
 
+  copyVar n = "$_copy_" <> n
+  copyFnArg a = a { name = copyVar a.name }
+
   collectFnArgs = go [] identity
     where
     go acc f (AST.Function fn@{ arguments, body: Just (AST.Block sts) })
@@ -70,9 +72,10 @@ tco = AST.everywhere convert
     go acc f (AST.Return (AST.Function fn@{ arguments, body: Just (AST.Block [ body ]) })) =
       go (map _.name arguments : acc) <@> body $ \b ->
         f $
-          AST.Function $ fn
-            { body = Just $ AST.Block [ b ]
-            }
+          AST.Return $
+            AST.Function $ fn
+              { body = Just $ AST.Block [ b ]
+              }
 
     go acc f (AST.Return (AST.Function fn@{ arguments, body: Just body@(AST.Block _) })) =
       (map _.name arguments : acc) /\ body /\ \b ->
@@ -121,6 +124,17 @@ tco = AST.everywhere convert
   toLoop :: String -> Array String -> AST -> AST
   toLoop ident args ast =
     AST.Block $
+      (args <#> \arg ->
+          AST.VariableIntroduction
+            { name: copyVar arg
+            , type: R.any
+            , qualifiers: []
+            , initialization:
+                Just $
+                  AST.App R.purs_any_copy
+                    [ AST.Var arg ]
+            })
+      <>
       [ AST.VariableIntroduction
           { name: tcoDone
           , type: R.any
@@ -145,23 +159,39 @@ tco = AST.everywhere convert
               Just $
                 AST.Function
                   { name: Just tcoLoop
-                  , arguments: []
+                  , arguments:
+                      [ { name: tcoDone
+                        , type: R.any
+                        }
+                      ] <> do
+                        args <#> \name ->
+                          { name: copyVar name
+                          , type: R.any
+                          }
                   , qualifiers: []
                   , returnType: R.any
                   , variadic:  false
                   , body:
                       Just $
-                        AST.Block
+                        AST.Block $
+                          (args <#> \arg ->
+                              AST.Assignment
+                                (AST.Var arg)
+                                (AST.App R.purs_any_copy
+                                  [ AST.Var $ copyVar arg ]))
+                          <>
                           [ loopify ast ]
                   }
           }
       , AST.While (AST.Unary AST.Not (AST.App R.purs_any_get_int [ AST.Var tcoDone ])) $
           AST.Block
             [ AST.Assignment (AST.Var tcoResult) $
                 AST.App R.purs_any_app $
-                  [ AST.Var tcoLoop
-                  , AST.Null
-                  ]
+                  A.concat $
+                    [ [ AST.Var tcoLoop, AST.Var tcoDone ]
+                    , AST.Var <<< copyVar <$> args
+                    , [ AST.Null ]
+                    ]
             ]
       , AST.Return $ AST.Var tcoResult
       ]
@@ -177,7 +207,7 @@ tco = AST.everywhere convert
           AST.Block $
             A.zipWith
               (\val arg ->
-                AST.App R.purs_any_assign_mut [ AST.Var arg, val ])
+                AST.App R.purs_any_assign_mut [ AST.Var $ copyVar arg, val ])
               allArgumentValues
               args
             <>

src/Language/PureScript/CodeGen/C/Pretty.purs

@@ -1,6 +1,7 @@
 module Language.PureScript.CodeGen.C.Pretty
   ( empty
   , prettyPrint
+  , renderType
   , PrintError(..)
   ) where
 
@@ -19,7 +20,6 @@ import Data.Maybe (Maybe(..))
 import Data.Newtype (unwrap)
 import Data.String.CodeUnits as CodeUnits
 import Data.Traversable (for_, traverse)
-import Debug.Trace (traceM)
 import Language.PureScript.CodeGen.C.AST (AST, PrimitiveType, Type, ValueQualifier)
 import Language.PureScript.CodeGen.C.AST as AST
 import Language.PureScript.CodeGen.C.AST as Type
@@ -123,7 +123,7 @@ prettyPrintAst x@(AST.Function
 
   let
     debugLambdas =
-      false
+      true
 
   name <-
     case debugLambdas, mName of

src/Language/PureScript/CodeGen/C/Transforms.purs

@@ -11,6 +11,7 @@ import Control.MonadPlus (guard)
 import Data.Array as A
 import Data.Either (Either(..))
 import Data.Function (on)
+import Data.Map as Map
 import Data.Maybe (Maybe(..), fromMaybe, maybe)
 import Data.Set as Set
 import Data.Traversable (for, traverse)
@@ -20,8 +21,10 @@ import Language.PureScript.CodeGen.C.AST (AST)
 import Language.PureScript.CodeGen.C.AST as AST
 import Language.PureScript.CodeGen.C.AST as Type
 import Language.PureScript.CodeGen.C.Common (freshName, isInternalVariable)
+import Language.PureScript.CodeGen.C.Pretty as PP
 import Language.PureScript.CodeGen.Runtime as R
 import Language.PureScript.CodeGen.SupplyT (class MonadSupply, freshId)
+import Math as Math
 
 -- | Split out variable declarations and definitions on a per-block (scope)
 -- | level and hoist the declarations to the top of the scope.
@@ -253,9 +256,10 @@ eraseLambdas moduleName asts =
             [ { name: "$_ctx"
               , type: Type.Pointer (Type.RawType scopeStruct.name [])
               }
-            ] <>
-            arguments <>
-            [ { name: "$_va_args"
+            , { name: fromMaybe "$_unused" $ _.name <$> A.head arguments
+              , type: R.any
+              }
+            , { name: "$_va_args"
               , type: Type.RawType "va_list" []
               }
             ]
@@ -265,19 +269,50 @@ eraseLambdas moduleName asts =
         , body:
             Just $
               AST.Block $
-                (A.catMaybes $
-                  A.mapWithIndex <@> scopeStruct.members $ \i varName -> ado
-                    guard $ not (A.elem varName $ _.name <$> arguments)
-                    in AST.VariableIntroduction
-                      { name: varName
-                      , type: R.any
-                      , qualifiers: []
-                      , initialization:
-                          Just $
-                            AST.Accessor (AST.Var varName) $
-                              AST.Var "$_ctx"
-                      }
-                ) <> [ body' ]
+                let
+                  bindings =
+                    Map.fromFoldable $
+                      A.filter
+                        (maybe
+                          (const true)
+                          (\x -> \v -> notEq (fst v) x)
+                          (_.name <$> A.head arguments)
+                        ) $
+                          A.concat $
+                            [ A.mapWithIndex <@> scopeStruct.members $
+                                \offset name ->
+                                  name /\ Just offset
+                            , (_ /\ Nothing) <<< _.name <$> do
+                                fromMaybe [] $ A.tail arguments
+                            ]
+                in
+                 A.concat
+                  [ Map.toUnfoldable bindings <#> \(name /\ mOffset) ->
+                      case mOffset of
+                        Nothing ->
+                          AST.VariableIntroduction
+                            { name
+                            , type: R.any
+                            , qualifiers: []
+                            , initialization:
+                                Just $
+                                  AST.App (AST.Var "va_arg")
+                                    [ AST.Var "$_va_args"
+                                    , AST.Raw $ PP.renderType R.any
+                                    ]
+                            }
+                        Just _ ->
+                          AST.VariableIntroduction
+                            { name
+                            , type: R.any
+                            , qualifiers: []
+                            , initialization:
+                                Just $
+                                  AST.Accessor (AST.Var name) $
+                                    AST.Var "$_ctx"
+                            }
+                  , [ body' ]
+                  ]
         }
 
     -- build up the continuation context and return it.

src/Language/PureScript/CodeGen/Runtime.purs

@@ -36,6 +36,7 @@ module Language.PureScript.CodeGen.Runtime
   , purs_any_num_new
   , purs_any_string_new
   , purs_any_char_new
+  , purs_any_copy
 
     -- code-gen helpers
   , _purs_scope_alloc
@@ -193,6 +194,9 @@ purs_any_cont_new = AST.Var "purs_any_cont_new"
 purs_any_string_new :: AST
 purs_any_string_new = AST.Var "purs_any_string_new"
 
+purs_any_copy :: AST
+purs_any_copy = AST.Var "purs_any_copy"
+
 purs_indirect_thunk_new :: AST
 purs_indirect_thunk_new = AST.Var "purs_indirect_thunk_new"
 

src/Main.purs

@@ -9,27 +9,27 @@ import Prelude
 import Control.Alt ((<|>))
 import Control.Monad.Error.Class (catchError, throwError)
 import Control.Monad.Except (except, runExcept, runExceptT, withExceptT)
-import CoreFn.FromJSON as C
-import CoreFn.Module as C
+import CoreFn.FromJSON (moduleFromJSON) as C
+import CoreFn.Module (FilePath(..), Module(..)) as C
 import Data.Array as A
 import Data.Either (Either(..), either)
 import Data.Maybe (Maybe(..), fromMaybe, maybe)
 import Data.Newtype (wrap)
 import Data.String as Str
-import Data.Traversable (for_, sequence, sequence_)
+import Data.Traversable (for_, sequence_)
 import Effect (Effect)
 import Effect.Aff (Aff, launchAff_)
 import Effect.Class (class MonadEffect, liftEffect)
 import Effect.Class.Console as Console
 import Effect.Exception (Error)
 import Effect.Exception as Error
-import Language.PureScript.CodeGen.C as C
+import Language.PureScript.CodeGen.C (moduleToAST) as C
 import Language.PureScript.CodeGen.C.AST as AST
 import Language.PureScript.CodeGen.C.File (dottedModuleName)
 import Language.PureScript.CodeGen.C.File as F
-import Language.PureScript.CodeGen.C.Pretty as C
+import Language.PureScript.CodeGen.C.Pretty (PrintError, prettyPrint) as C
 import Language.PureScript.CodeGen.C.Pretty as PrintError
-import Language.PureScript.CodeGen.CompileError as C
+import Language.PureScript.CodeGen.CompileError (CompileError) as C
 import Language.PureScript.CodeGen.CompileError as CompileError
 import Language.PureScript.CodeGen.SupplyT (runSupplyT)
 import Node.Encoding (Encoding(..))