Move outermost lambdas to top level

This improves runtime performance a bit since it avoids creating some
unnecessary closures (and is a preliminary step to lambda lifting, which
might help for some back ends).

src/Compiler/Common.idr

@@ -101,12 +101,10 @@ fastAppend xs
     build b (x :: xs) = do addToStringBuffer b x
                            build b xs
 
-dumpCases : {auto c : Ref Ctxt Defs} ->
-            String -> List Name ->
+dumpCases : Defs -> String -> List Name ->
             Core ()
-dumpCases fn cns
-    = do defs <- get Ctxt
-         cstrs <- traverse (dumpCase defs) cns
+dumpCases defs fn cns
+    = do cstrs <- traverse dumpCase cns
          Right () <- coreLift $ writeFile fn (fastAppend cstrs)
                | Left err => throw (FileErr fn err)
          pure ()
@@ -115,8 +113,8 @@ dumpCases fn cns
     fullShow (DN _ n) = show n
     fullShow n = show n
 
-    dumpCase : Defs -> Name -> Core String
-    dumpCase defs n
+    dumpCase : Name -> Core String
+    dumpCase n
         = case !(lookupCtxtExact n (gamma defs)) of
                Nothing => pure ""
                Just d =>
@@ -125,10 +123,11 @@ dumpCases fn cns
                          Just def => pure (fullShow n ++ " = " ++ show def ++ "\n")
 
 -- Find all the names which need compiling, from a given expression, and compile
--- them to CExp form (and update that in the Defs)
+-- them to CExp form (and update that in the Defs).
+-- Return the names, the type tags, and a compiled version of the expression
 export
 findUsedNames : {auto c : Ref Ctxt Defs} ->
-                Term vars -> Core (List Name, NameTags)
+                ClosedTerm -> Core (List Name, NameTags, NamedCExp)
 findUsedNames tm
     = do defs <- get Ctxt
          sopts <- getSession
@@ -154,12 +153,17 @@ findUsedNames tm
          logTime ("Compile defs " ++ show (length cns) ++ "/" ++ show asize) $
            traverse_ (compileDef tycontags) cns
          logTime "Inline" $ traverse_ inlineDef cns
+         logTime "Merge lambda" $ traverse_ mergeLamDef cns
+         logTime "Fix arity" $ traverse_ fixArityDef cns
          logTime "Forget names" $ traverse_ mkForgetDef cns
+         let compiledtm = forget !(fixArityExp !(compileExp tycontags tm))
+
+         defs <- get Ctxt
          maybe (pure ())
                (\f => do coreLift $ putStrLn $ "Dumping case trees to " ++ f
-                         dumpCases f cns)
+                         dumpCases defs f cns)
                (dumpcases sopts)
-         pure (cns, tycontags)
+         pure (cns, tycontags, compiledtm)
   where
     primTags : Int -> NameTags -> List Constant -> NameTags
     primTags t tags [] = tags

src/Compiler/CompileExpr.idr

@@ -60,6 +60,7 @@ etaExpand i (S k) exp args
              (etaExpand (i + 1) k (weaken exp)
                   (MkVar First :: map weakenVar args))
 
+export
 expandToArity : Nat -> CExp vars -> List (CExp vars) -> CExp vars
 -- No point in applying to anything
 expandToArity _ (CErased fc) _ = CErased fc
@@ -472,10 +473,10 @@ toCDef tags n ty def
 
 export
 compileExp : {auto c : Ref Ctxt Defs} ->
-             NameTags -> ClosedTerm -> Core NamedCExp
+             NameTags -> ClosedTerm -> Core (CExp [])
 compileExp tags tm
     = do exp <- toCExp tags (UN "main") tm
-         pure (forget exp)
+         pure exp
 
 ||| Given a name, look up an expression, and compile it to a CExp in the environment
 export

src/Compiler/Inline.idr

@@ -1,5 +1,7 @@
 module Compiler.Inline
 
+import Compiler.CompileExpr
+
 import Core.CompileExpr
 import Core.Context
 import Core.FC
@@ -298,12 +300,124 @@ mutual
            else pickConstAlt rec env stk (CPrimVal fc c) alts def
   pickConstAlt rec env stk _ _ _ = pure Nothing
 
-inline : {auto c : Ref Ctxt Defs} ->
-         CDef -> Core CDef
-inline (MkFun args exp)
+-- Inlining may have messed with function arity (e.g. by adding lambdas to
+-- the LHS to avoid needlessly making a closure) so fix them up here. This
+-- needs to be right because typically back ends need to know whether a
+-- name is under- or over-applied
+fixArityTm : {auto c : Ref Ctxt Defs} ->
+             CExp vars -> List (CExp vars) -> Core (CExp vars)
+fixArityTm (CRef fc n) args
+    = do defs <- get Ctxt
+         Just gdef <- lookupCtxtExact n (gamma defs)
+              | Nothing => pure (unload args (CRef fc n))
+         let Just def = compexpr gdef
+              | Nothing => pure (unload args (CRef fc n))
+         let arity = getArity def
+         pure $ expandToArity arity (CApp fc (CRef fc n) []) args
+fixArityTm (CLam fc x sc) args
+    = pure $ expandToArity Z (CLam fc x !(fixArityTm sc [])) args
+fixArityTm (CLet fc x val sc) args
+    = pure $ expandToArity Z
+                 (CLet fc x !(fixArityTm val []) !(fixArityTm sc [])) args
+fixArityTm (CApp fc f fargs) args
+    = fixArityTm f (!(traverse (\tm => fixArityTm tm []) fargs) ++ args)
+fixArityTm (CCon fc n t args) []
+    = pure $ CCon fc n t !(traverse (\tm => fixArityTm tm []) args)
+fixArityTm (COp fc op args) []
+    = pure $ COp fc op !(traverseArgs args)
+  where
+    traverseArgs : Vect n (CExp vs) -> Core (Vect n (CExp vs))
+    traverseArgs [] = pure []
+    traverseArgs (a :: as) = pure $ !(fixArityTm a []) :: !(traverseArgs as)
+fixArityTm (CExtPrim fc p args) []
+    = pure $ CExtPrim fc p !(traverse (\tm => fixArityTm tm []) args)
+fixArityTm (CForce fc tm) args
+    = pure $ expandToArity Z (CForce fc !(fixArityTm tm [])) args
+fixArityTm (CDelay fc tm) args
+    = pure $ expandToArity Z (CDelay fc !(fixArityTm tm [])) args
+fixArityTm (CConCase fc sc alts def) args
+    = pure $ expandToArity Z
+              (CConCase fc !(fixArityTm sc [])
+                           !(traverse fixArityAlt alts)
+                           !(traverseOpt (\tm => fixArityTm tm []) def)) args
+  where
+    fixArityAlt : CConAlt vars -> Core (CConAlt vars)
+    fixArityAlt (MkConAlt n t a sc)
+        = pure $ MkConAlt n t a !(fixArityTm sc [])
+fixArityTm (CConstCase fc sc alts def) args
+    = pure $ expandToArity Z
+              (CConstCase fc !(fixArityTm sc [])
+                             !(traverse fixArityConstAlt alts)
+                             !(traverseOpt (\tm => fixArityTm tm []) def)) args
+  where
+    fixArityConstAlt : CConstAlt vars -> Core (CConstAlt vars)
+    fixArityConstAlt (MkConstAlt c sc)
+        = pure $ MkConstAlt c !(fixArityTm sc [])
+fixArityTm t [] = pure t
+fixArityTm t args = pure $ expandToArity Z t args
+
+export
+fixArityExp : {auto c : Ref Ctxt Defs} ->
+              CExp vars -> Core (CExp vars)
+fixArityExp tm = fixArityTm tm []
+
+fixArity : {auto c : Ref Ctxt Defs} ->
+           CDef -> Core CDef
+fixArity (MkFun args exp) = pure $ MkFun args !(fixArityTm exp [])
+fixArity (MkError exp) = pure $ MkError !(fixArityTm exp [])
+fixArity d = pure d
+
+getLams : Int -> SubstCEnv done args -> CExp (done ++ args) ->
+          (args' ** (SubstCEnv args' args, CExp (args' ++ args)))
+getLams {done} i env (CLam fc x sc)
+    = getLams {done = x :: done} (i + 1) (CRef fc (MN "ext" i) :: env) sc
+getLams {done} i env sc = (done ** (env, sc))
+
+mkBounds : (xs : _) -> Bounds xs
+mkBounds [] = None
+mkBounds (x :: xs) = Add x x (mkBounds xs)
+
+getNewArgs : SubstCEnv done args -> List Name
+getNewArgs [] = []
+getNewArgs (CRef _ n :: xs) = n :: getNewArgs xs
+getNewArgs {done = x :: xs} (_ :: sub) = x :: getNewArgs sub
+
+-- Move any lambdas in the body of the definition into the lhs list of vars.
+-- Annoyingly, the indices will need fixing up because the order in the top
+-- level definition goes left to right (i.e. first argument has lowest index,
+-- not the highest, as you'd expect if they were all lambdas).
+mergeLambdas : (args : List Name) -> CExp args -> (args' ** CExp args')
+mergeLambdas args (CLam fc x sc)
+    = let (args' ** (env, exp')) = getLams 0 [] (CLam fc x sc)
+          expNs = substs env exp'
+          newArgs = reverse $ getNewArgs env
+          expLocs = mkLocals {later = args} {vars=[]} (mkBounds newArgs)
+                             (rewrite appendNilRightNeutral args in expNs) in
+          (_ ** expLocs)
+mergeLambdas args exp = (args ** exp)
+
+doEval : {auto c : Ref Ctxt Defs} ->
+         Name -> CExp args -> Core (CExp args)
+doEval n exp
     = do l <- newRef LVar (the Int 0)
-         MkFun args <$> eval [] [] [] exp
-inline d = pure d
+         log 10 (show n ++ ": " ++ show exp)
+         exp' <- eval [] [] [] exp
+         log 10 ("Inlined: " ++ show exp')
+         pure exp'
+
+inline : {auto c : Ref Ctxt Defs} ->
+         Name -> CDef -> Core CDef
+inline n (MkFun args def)
+    = pure $ MkFun args !(doEval n def)
+inline n d = pure d
+
+-- merge lambdas from expression into top level arguments
+mergeLam : {auto c : Ref Ctxt Defs} ->
+           CDef -> Core CDef
+mergeLam (MkFun args def)
+    = do let (args' ** exp') = mergeLambdas args def
+         pure $ MkFun args' exp'
+mergeLam d = pure d
 
 export
 inlineDef : {auto c : Ref Ctxt Defs} ->
@@ -312,5 +426,22 @@ inlineDef n
     = do defs <- get Ctxt
          Just def <- lookupCtxtExact n (gamma defs) | Nothing => pure ()
          let Just cexpr =  compexpr def             | Nothing => pure ()
-         setCompiled n !(inline cexpr)
+         setCompiled n !(inline n cexpr)
 
+export
+fixArityDef : {auto c : Ref Ctxt Defs} ->
+              Name -> Core ()
+fixArityDef n
+    = do defs <- get Ctxt
+         Just def <- lookupCtxtExact n (gamma defs) | Nothing => pure ()
+         let Just cexpr =  compexpr def             | Nothing => pure ()
+         setCompiled n !(fixArity cexpr)
+
+export
+mergeLamDef : {auto c : Ref Ctxt Defs} ->
+              Name -> Core ()
+mergeLamDef n
+    = do defs <- get Ctxt
+         Just def <- lookupCtxtExact n (gamma defs) | Nothing => pure ()
+         let Just cexpr =  compexpr def             | Nothing => pure ()
+         setCompiled n !(mergeLam cexpr)

src/Compiler/Scheme/Chez.idr

@@ -330,14 +330,14 @@ compileToSS c tm outfile
     = do ds <- getDirectives Chez
          libs <- findLibs ds
          traverse_ copyLib libs
-         (ns, tags) <- findUsedNames tm
+         (ns, tags, ctm) <- findUsedNames tm
          defs <- get Ctxt
          l <- newRef {t = List String} Loaded ["libc", "libc 6"]
          s <- newRef {t = List String} Structs []
          fgndefs <- traverse getFgnCall ns
          compdefs <- traverse (getScheme chezExtPrim chezString defs) ns
          let code = fastAppend (map snd fgndefs ++ compdefs)
-         main <- schExp chezExtPrim chezString 0 !(compileExp tags tm)
+         main <- schExp chezExtPrim chezString 0 ctm
          chez <- coreLift findChez
          support <- readDataFile "chez/support.ss"
          let scm = schHeader chez (map snd libs) ++

src/Compiler/Scheme/Racket.idr

@@ -291,14 +291,14 @@ getFgnCall n
 compileToRKT : Ref Ctxt Defs ->
                ClosedTerm -> (outfile : String) -> Core ()
 compileToRKT c tm outfile
-    = do (ns, tags) <- findUsedNames tm
+    = do (ns, tags, ctm) <- findUsedNames tm
          defs <- get Ctxt
          l <- newRef {t = List String} Loaded []
          s <- newRef {t = List String} Structs []
          fgndefs <- traverse getFgnCall ns
          compdefs <- traverse (getScheme racketPrim racketString defs) ns
          let code = fastAppend (map snd fgndefs ++ compdefs)
-         main <- schExp racketPrim racketString 0 !(compileExp tags tm)
+         main <- schExp racketPrim racketString 0 ctm
          support <- readDataFile "racket/support.rkt"
          let scm = schHeader (concat (map fst fgndefs)) ++
                    support ++ code ++