Reorganise TTC to make code generation faster

Before compilation we need to decode all the relevant definitions from
the TTC, but the only bit we actually need is the compiled IR. So, put
that first, and leave a dummy definition for the rest, then replace it
when we're done.
Chasing all the definitions now takes about 20% of the time it did
before (at least on the program I'm most interested in)

src/Compiler/Common.idr

@@ -10,6 +10,7 @@ import Core.Context
 import Core.Directory
 import Core.Options
 import Core.TT
+import Core.TTC
 import Utils.Binary
 
 import Data.IOArray
@@ -29,12 +30,33 @@ record Codegen where
   ||| Execute an Idris 2 expression directly.
   executeExpr : Ref Ctxt Defs -> (execDir : String) -> ClosedTerm -> Core ()
 
+-- Say which phase of compilation is the last one to use - it saves time if 
+-- you only ask for what you need.
+public export
+data UsePhase = Cases | Lifted | ANF | VMCode
+
+Eq UsePhase where
+  (==) Cases Cases = True
+  (==) Lifted Lifted = True
+  (==) ANF ANF = True
+  (==) VMCode VMCode = True
+  (==) _ _ = False
+
+Ord UsePhase where
+  compare x y = compare (tag x) (tag y)
+    where
+      tag : UsePhase -> Int
+      tag Cases = 0
+      tag Lifted = 0
+      tag ANF = 0
+      tag VMCode = 0
+
 public export
 record CompileData where
   constructor MkCompileData
-  allNames : List Name -- names which need to be compiled
   mainExpr : CExp [] -- main expression to execute. This also appears in
                      -- the definitions below as MN "__mainExpression" 0
+  namedDefs : List (Name, FC, NamedDef)
   lambdaLifted : List (Name, LiftedDef)
        -- ^ lambda lifted definitions, if required. Only the top level names
        -- will be in the context, and (for the moment...) I don't expect to
@@ -56,7 +78,8 @@ compile : {auto c : Ref Ctxt Defs} ->
 compile {c} cg tm out
     = do makeExecDirectory
          d <- getDirs
-         compileExpr cg c (exec_dir d) tm out
+         logTime "Code generation overall" $
+             compileExpr cg c (exec_dir d) tm out
 
 ||| execute
 ||| As with `compile`, produce a functon that executes
@@ -70,27 +93,73 @@ execute {c} cg tm
          executeExpr cg c (exec_dir d) tm
          pure ()
 
+-- If an entry isn't already decoded, get the minimal entry we need for
+-- compilation, and record the Binary so that we can put it back when we're
+-- done (so that we don't obliterate the definition)
+getMinimalDef : ContextEntry -> Core (GlobalDef, Maybe Binary)
+getMinimalDef (Decoded def) = pure (def, Nothing)
+getMinimalDef (Coded bin)
+    = do b <- newRef Bin bin
+         cdef <- fromBuf b
+         refsRList <- fromBuf b
+         let refsR = map fromList refsRList
+         fc <- fromBuf b
+         mul <- fromBuf b
+         name <- fromBuf b
+         let def
+             = MkGlobalDef fc name (Erased fc False) [] [] [] [] mul
+                           [] Public (MkTotality Unchecked IsCovering)
+                           [] Nothing refsR False False True
+                           None cdef Nothing []
+         pure (def, Just bin)
+
 -- ||| Recursively get all calls in a function definition
 getAllDesc : {auto c : Ref Ctxt Defs} ->
              List Name -> -- calls to check
-             IOArray Int -> -- which nodes have been visited. If the entry is
-                            -- present, it's visited
+             IOArray (Int, Maybe Binary) ->
+                            -- which nodes have been visited. If the entry is
+                            -- present, it's visited. Keep the binary entry, if
+                            -- we partially decoded it, so that we can put back
+                            -- the full definition later.
+                            -- (We only need to decode the case tree IR, and
+                            -- it's expensive to decode the whole thing)
              Defs -> Core ()
 getAllDesc [] arr defs = pure ()
 getAllDesc (n@(Resolved i) :: rest) arr defs
   = do Nothing <- coreLift $ readArray arr i
            | Just _ => getAllDesc rest arr defs
-       case !(lookupCtxtExact n (gamma defs)) of
+       case !(lookupContextEntry n (gamma defs)) of
             Nothing => getAllDesc rest arr defs
-            Just def =>
-              if multiplicity def /= erased
-                 then do coreLift $ writeArray arr i i
-                         let refs = refersToRuntime def
-                         getAllDesc (keys refs ++ rest) arr defs
-                 else getAllDesc rest arr defs
+            Just (_, entry) =>
+              do (def, bin) <- getMinimalDef entry
+                 addDef n def 
+                 let refs = refersToRuntime def
+                 if multiplicity def /= erased
+                    then do coreLift $ writeArray arr i (i, bin)
+                            let refs = refersToRuntime def
+                            refs' <- traverse toResolvedNames (keys refs)
+                            getAllDesc (refs' ++ rest) arr defs
+                    else getAllDesc rest arr defs
 getAllDesc (n :: rest) arr defs
   = getAllDesc rest arr defs
 
+getNamedDef : {auto c : Ref Ctxt Defs} ->
+              Name -> Core (Maybe (Name, FC, NamedDef))
+getNamedDef n
+    = do defs <- get Ctxt
+         case !(lookupCtxtExact n (gamma defs)) of
+              Nothing => pure Nothing
+              Just def => case namedcompexpr def of
+                               Nothing => pure Nothing
+                               Just d => pure (Just (n, location def, d))
+
+replaceEntry : {auto c : Ref Ctxt Defs} ->
+               (Int, Maybe Binary) -> Core ()
+replaceEntry (i, Nothing) = pure ()
+replaceEntry (i, Just b)
+    = do addContextEntry (Resolved i) b
+         pure ()
+
 natHackNames : List Name
 natHackNames
     = [UN "prim__add_Integer",
@@ -182,8 +251,8 @@ dumpVMCode fn lns
 -- Return the names, the type tags, and a compiled version of the expression
 export
 getCompileData : {auto c : Ref Ctxt Defs} ->
-                 ClosedTerm -> Core CompileData
-getCompileData tm_in
+                 UsePhase -> ClosedTerm -> Core CompileData
+getCompileData phase tm_in
     = do defs <- get Ctxt
          sopts <- getSession
          let ns = getRefs (Resolved (-1)) tm_in
@@ -194,8 +263,9 @@ getCompileData tm_in
          asize <- getNextEntry
          arr <- coreLift $ newArray asize
          logTime "Get names" $ getAllDesc (natHackNames' ++ keys ns) arr defs
-         let allNs = mapMaybe (maybe Nothing (Just . Resolved))
-                              !(coreLift (toList arr))
+
+         let entries = mapMaybe id !(coreLift (toList arr))
+         let allNs = map (Resolved . fst) entries
          cns <- traverse toFullNames allNs
 
          -- Do a round of merging/arity fixing for any names which were
@@ -210,13 +280,20 @@ getCompileData tm_in
          let mainname = MN "__mainExpression" 0
          (liftedtm, ldefs) <- liftBody mainname compiledtm
 
-         lifted_in <- logTime "Lambda lift" $ traverse lambdaLift rcns
+         namedefs <- traverse getNamedDef rcns
+         lifted_in <- if phase >= Lifted
+                         then logTime "Lambda lift" $ traverse lambdaLift rcns
+                         else pure []
 
          let lifted = (mainname, MkLFun [] [] liftedtm) ::
                       ldefs ++ concat lifted_in
 
-         anf <- logTime "Get ANF" $ traverse (\ (n, d) => pure (n, !(toANF d))) lifted
-         vmcode <- logTime "Get VM Code" $ pure (allDefs anf)
+         anf <- if phase >= ANF
+                   then logTime "Get ANF" $ traverse (\ (n, d) => pure (n, !(toANF d))) lifted
+                   else pure []
+         vmcode <- if phase >= VMCode
+                      then logTime "Get VM Code" $ pure (allDefs anf)
+                      else pure []
 
          defs <- get Ctxt
          maybe (pure ())
@@ -235,7 +312,13 @@ getCompileData tm_in
                (\f => do coreLift $ putStrLn $ "Dumping VM defs to " ++ f
                          dumpVMCode f vmcode)
                (dumpvmcode sopts)
-         pure (MkCompileData rcns compiledtm
+
+         -- We're done with our minimal context now, so put it back the way
+         -- it was. Back ends shouldn't look at the global context, because
+         -- it'll have to decode the definitions again.
+         traverse_ replaceEntry entries
+         pure (MkCompileData compiledtm
+                             (mapMaybe id namedefs)
                              lifted anf vmcode)
   where
     nonErased : Name -> Core Bool

src/Compiler/Scheme/Chez.idr

@@ -282,8 +282,8 @@ mkStruct _ = pure ""
 schFgnDef : {auto c : Ref Ctxt Defs} ->
             {auto l : Ref Loaded (List String)} ->
             {auto s : Ref Structs (List String)} ->
-            String -> FC -> Name -> CDef -> Core (String, String)
-schFgnDef appdir fc n (MkForeign cs args ret)
+            String -> FC -> Name -> NamedDef -> Core (String, String)
+schFgnDef appdir fc n (MkNmForeign cs args ret)
     = do let argns = mkArgs 0 args
          let allargns = map fst argns
          let useargns = map fst (filter snd argns)
@@ -301,15 +301,8 @@ schFgnDef _ _ _ _ = pure ("", "")
 getFgnCall : {auto c : Ref Ctxt Defs} ->
              {auto l : Ref Loaded (List String)} ->
              {auto s : Ref Structs (List String)} ->
-             String -> Name -> Core (String, String)
-getFgnCall appdir n
-    = do defs <- get Ctxt
-         case !(lookupCtxtExact n (gamma defs)) of
-           Nothing => throw (InternalError ("Compiling undefined name " ++ show n))
-           Just def => case compexpr def of
-                          Nothing =>
-                             throw (InternalError ("No compiled definition for " ++ show n))
-                          Just d => schFgnDef appdir (location def) n d
+             String -> (Name, FC, NamedDef) -> Core (String, String)
+getFgnCall appdir (n, fc, d) = schFgnDef appdir fc n d
 
 startChez : String -> String
 startChez target = unlines
@@ -326,15 +319,15 @@ compileToSS c appdir tm outfile
     = do ds <- getDirectives Chez
          libs <- findLibs ds
          traverse_ copyLib libs
-         cdata <- getCompileData tm
-         let ns = allNames cdata
+         cdata <- getCompileData Cases tm
+         let ndefs = namedDefs cdata
          let ctm = forget (mainExpr cdata)
 
          defs <- get Ctxt
          l <- newRef {t = List String} Loaded ["libc", "libc 6"]
          s <- newRef {t = List String} Structs []
-         fgndefs <- traverse (getFgnCall appdir) ns
-         compdefs <- traverse (getScheme chezExtPrim chezString defs) ns
+         fgndefs <- traverse (getFgnCall appdir) ndefs
+         compdefs <- traverse (getScheme chezExtPrim chezString) ndefs
          let code = fastAppend (map snd fgndefs ++ compdefs)
          main <- schExp chezExtPrim chezString 0 ctm
          chez <- coreLift findChez

src/Compiler/Scheme/Common.idr

@@ -448,11 +448,6 @@ export
 getScheme : {auto c : Ref Ctxt Defs} ->
             (schExtPrim : Int -> ExtPrim -> List NamedCExp -> Core String) ->
             (schString : String -> String) ->
-            Defs -> Name -> Core String
-getScheme schExtPrim schString defs n
-    = case !(lookupCtxtExact n (gamma defs)) of
-           Nothing => throw (InternalError ("Compiling undefined name " ++ show n))
-           Just d => case namedcompexpr d of
-                          Nothing =>
-                             throw (InternalError ("No compiled definition for " ++ show n))
-                          Just d => schDef schExtPrim schString n d
+            (Name, FC, NamedDef) -> Core String
+getScheme schExtPrim schString (n, fc, d)
+    = schDef schExtPrim schString n d

src/Compiler/Scheme/Gambit.idr

@@ -245,8 +245,8 @@ mkStruct _ = pure ""
 schFgnDef : {auto c : Ref Ctxt Defs} ->
             {auto l : Ref Loaded (List String)} ->
             {auto s : Ref Structs (List String)} ->
-            FC -> Name -> CDef -> Core (String, String)
-schFgnDef fc n (MkForeign cs args ret)
+            FC -> Name -> NamedDef -> Core (String, String)
+schFgnDef fc n (MkNmForeign cs args ret)
     = do let argns = mkArgs 0 args
          let allargns = map fst argns
          let useargns = map fst (filter snd argns)
@@ -264,30 +264,23 @@ schFgnDef _ _ _ = pure ("", "")
 getFgnCall : {auto c : Ref Ctxt Defs} ->
              {auto l : Ref Loaded (List String)} ->
              {auto s : Ref Structs (List String)} ->
-             Name -> Core (String, String)
-getFgnCall n
-    = do defs <- get Ctxt
-         case !(lookupCtxtExact n (gamma defs)) of
-           Nothing => throw (InternalError ("Compiling undefined name " ++ show n))
-           Just def => case compexpr def of
-                          Nothing =>
-                             throw (InternalError ("No compiled definition for " ++ show n))
-                          Just d => schFgnDef (location def) n d
+             (Name, FC, NamedDef) -> Core (String, String)
+getFgnCall (n, fc, d) = schFgnDef fc n d
 
 -- TODO Include libraries from the directives
 compileToSCM : Ref Ctxt Defs ->
                ClosedTerm -> (outfile : String) -> Core ()
 compileToSCM c tm outfile
-    = do cdata <- getCompileData tm
-         let ns = allNames cdata
+    = do cdata <- getCompileData Cases tm
+         let ndefs = namedDefs cdata
          -- let tags = nameTags cdata
          let ctm = forget (mainExpr cdata)
 
          defs <- get Ctxt
          l <- newRef {t = List String} Loaded []
          s <- newRef {t = List String} Structs []
-         fgndefs <- traverse getFgnCall ns
-         compdefs <- traverse (getScheme gambitPrim gambitString defs) ns
+         fgndefs <- traverse getFgnCall ndefs
+         compdefs <- traverse (getScheme gambitPrim gambitString) ndefs
          let code = fastAppend (map snd fgndefs ++ compdefs) ++
                     concat (map fst fgndefs)
          main <- schExp gambitPrim gambitString 0 ctm

src/Compiler/Scheme/Racket.idr

@@ -258,8 +258,8 @@ mkStruct _ = pure ""
 schFgnDef : {auto c : Ref Ctxt Defs} ->
             {auto l : Ref Loaded (List String)} ->
             {auto s : Ref Structs (List String)} ->
-            FC -> Name -> CDef -> Core (String, String)
-schFgnDef fc n (MkForeign cs args ret)
+            FC -> Name -> NamedDef -> Core (String, String)
+schFgnDef fc n (MkNmForeign cs args ret)
     = do let argns = mkArgs 0 args
          let allargns = map fst argns
          let useargns = map fst (filter snd argns)
@@ -276,28 +276,21 @@ schFgnDef _ _ _ = pure ("", "")
 getFgnCall : {auto c : Ref Ctxt Defs} ->
              {auto l : Ref Loaded (List String)} ->
              {auto s : Ref Structs (List String)} ->
-             Name -> Core (String, String)
-getFgnCall n
-    = do defs <- get Ctxt
-         case !(lookupCtxtExact n (gamma defs)) of
-           Nothing => throw (InternalError ("Compiling undefined name " ++ show n))
-           Just def => case compexpr def of
-                          Nothing =>
-                             throw (InternalError ("No compiled definition for " ++ show n))
-                          Just d => schFgnDef (location def) n d
+             (Name, FC, NamedDef) -> Core (String, String)
+getFgnCall (n, fc, d) = schFgnDef fc n d
 
 compileToRKT : Ref Ctxt Defs ->
                ClosedTerm -> (outfile : String) -> Core ()
 compileToRKT c tm outfile
-    = do cdata <- getCompileData tm
-         let ns = allNames cdata
+    = do cdata <- getCompileData Cases tm
+         let ndefs = namedDefs cdata
          let ctm = forget (mainExpr cdata)
 
          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
+         fgndefs <- traverse getFgnCall ndefs
+         compdefs <- traverse (getScheme racketPrim racketString) ndefs
          let code = fastAppend (map snd fgndefs ++ compdefs)
          main <- schExp racketPrim racketString 0 ctm
          support <- readDataFile "racket/support.rkt"

src/Core/Binary.idr

@@ -27,7 +27,7 @@ import Data.Buffer
 -- TTC files can only be compatible if the version number is the same
 export
 ttcVersion : Int
-ttcVersion = 25
+ttcVersion = 26
 
 export
 checkTTCVersion : String -> Int -> Int -> Core ()

src/Core/CompileExpr.idr

@@ -6,6 +6,7 @@ import Core.FC
 import Core.Name
 import Core.TT
 
+import Data.NameMap
 import Data.Vect
 
 %default covering

src/Core/Context.idr

@@ -246,7 +246,7 @@ export
 data Arr : Type where
 
 -- A context entry. If it's never been looked up, we haven't decoded the
--- binary blod yet, so decode it first time
+-- binary blob yet, so decode it first time
 public export
 data ContextEntry : Type where
      Coded : Binary -> ContextEntry