[ performance ] Implement weak memoisation of lazy values for chez an…

…d racket (#2791)

* [ lazy ] Weakly memoise lazy expressions on chez and racket

* [ lazy ] Make weak memoisation to be controlled by codegen directive

CHANGELOG_NEXT.md

@@ -129,12 +129,20 @@ This CHANGELOG describes the merged but unreleased changes. Please see [CHANGELO
 
 * More efficient `collect-request-handler` is used.
 
+* Add a codegen directive called `lazy=weakMemo` to make `Lazy` and `Inf` values *weakly*
+  memoised. That is, once accessed, they are allowed to be not re-evaluated until garbage
+  collector wipes them.
+
 #### Racket
 
 * Fixed CSE soundness bug that caused delayed expressions to sometimes be eagerly
   evaluated. Now when a delayed expression is lifted by CSE, it is compiled
   using Scheme's `delay` and `force` to memoize them.
 
+* Add a codegen directive called `lazy=weakMemo` to make `Lazy` and `Inf` values *weakly*
+  memoised. That is, once accessed, they are allowed to be not re-evaluated until garbage
+  collector wipes them.
+
 #### NodeJS Backend
 
 * The NodeJS executable output to `build/exec/` now has its executable bit set.

docs/source/backends/chez.rst

@@ -72,6 +72,15 @@ Chez Directives
 
     $ idris2 --codegen chez --directive extraRuntime=/path/to/extensions.scm -o main Main.idr
 
+* ``--directive lazy=weakMemo``
+
+  Makes all non-toplevel ``Lazy`` and ``Inf`` values to be *weakly* memoised.
+  That is, once this expression is evaluated at runtime, it is allowed to not to be recalculated on later accesses
+  until memoised value is wiped by a garbage collector.
+  Garbage collector is allowed to collect weakly memoised values at its own discretion,
+  so when no free memory is available, weakly memoised values are free to be wiped.
+  That's why it is safer comparing to full memoisation.
+
 Making a freestanding executable
 ================================
 

docs/source/backends/racket.rst

@@ -69,3 +69,12 @@ Racket Directives
   .. code-block::
 
     $ idris2 --codegen chez --directive extraRuntime=/path/to/extensions.scm -o main Main.idr
+
+* ``--directive lazy=weakMemo``
+
+  Makes all non-toplevel ``Lazy`` and ``Inf`` values to be *weakly* memoised.
+  That is, once this expression is evaluated at runtime, it is allowed to not to be recalculated on later accesses
+  until memoised value is wiped by a garbage collector.
+  Garbage collector is allowed to collect weakly memoised values at its own discretion,
+  so when no free memory is available, weakly memoised values are free to be wiped.
+  That's why it is safer comparing to full memoisation.

src/Compiler/Common.idr

@@ -557,6 +557,14 @@ getExtraRuntime directives
     paths : List String
     paths = nub $ mapMaybe getArg $ reverse directives
 
+-- parses `--directive lazy=weakMemo` option for turning on weak memoisation of lazy values
+-- (if supported by a backend).
+-- This particular form of the directive string is chosen to be able to pass different variants
+-- in the future (say, for strong memoisation, or turning laziness off).
+export
+getWeakMemoLazy : List String -> Bool
+getWeakMemoLazy = elem "lazy=weakMemo"
+
 ||| Cast implementations. Values of `ConstantPrimitives` can
 ||| be used in a call to `castInt`, which then determines
 ||| the cast implementation based on the given pair of

src/Compiler/Scheme/Chez.idr

@@ -133,45 +133,44 @@ export
 chezString : String -> Builder
 chezString cs = "\"" ++ showChezString (unpack cs) "\""
 
-mutual
-  handleRet : CFType -> Builder -> Builder
-  handleRet CFUnit op = op ++ " " ++ mkWorld (schConstructor chezString (UN Underscore) (Just 0) [])
-  handleRet _ op = mkWorld op
-
-  getFArgs : NamedCExp -> Core (List (NamedCExp, NamedCExp))
-  getFArgs (NmCon fc _ _ (Just 0) _) = pure []
-  getFArgs (NmCon fc _ _ (Just 1) [ty, val, rest]) = pure $ (ty, val) :: !(getFArgs rest)
-  getFArgs arg = throw (GenericMsg (getFC arg) ("Badly formed c call argument list " ++ show arg))
-
-  export
-  chezExtPrim : SortedSet Name -> Nat -> ExtPrim -> List NamedCExp -> Core Builder
-  chezExtPrim cs i GetField [NmPrimVal _ (Str s), _, _, struct,
-                          NmPrimVal _ (Str fld), _]
-      = do structsc <- schExp cs (chezExtPrim cs) chezString 0 struct
-           pure $ "(ftype-ref " ++ fromString s ++ " (" ++ fromString fld ++ ") " ++ structsc ++ ")"
-  chezExtPrim cs i GetField [_,_,_,_,_,_]
-      = pure "(blodwen-error-quit \"bad getField\")"
-  chezExtPrim cs i SetField [NmPrimVal _ (Str s), _, _, struct,
-                          NmPrimVal _ (Str fld), _, val, world]
-      = do structsc <- schExp cs (chezExtPrim cs) chezString 0 struct
-           valsc <- schExp cs (chezExtPrim cs) chezString 0 val
-           pure $ mkWorld $
-              "(ftype-set! " ++ fromString s ++ " (" ++ fromString fld ++ ") " ++ structsc ++
-              " " ++ valsc ++ ")"
-  chezExtPrim cs i SetField [_,_,_,_,_,_,_,_]
-      = pure "(blodwen-error-quit \"bad setField\")"
-  chezExtPrim cs i SysCodegen []
-      = pure $ "\"chez\""
-  chezExtPrim cs i OnCollect [_, p, c, world]
-      = do p' <- schExp cs (chezExtPrim cs) chezString 0 p
-           c' <- schExp cs (chezExtPrim cs) chezString 0 c
-           pure $ mkWorld $ "(blodwen-register-object " ++ p' ++ " " ++ c' ++ ")"
-  chezExtPrim cs i OnCollectAny [p, c, world]
-      = do p' <- schExp cs (chezExtPrim cs) chezString 0 p
-           c' <- schExp cs (chezExtPrim cs) chezString 0 c
-           pure $ mkWorld $ "(blodwen-register-object " ++ p' ++ " " ++ c' ++ ")"
-  chezExtPrim cs i prim args
-      = schExtCommon cs (chezExtPrim cs) chezString i prim args
+handleRet : CFType -> Builder -> Builder
+handleRet CFUnit op = op ++ " " ++ mkWorld (schConstructor chezString (UN Underscore) (Just 0) [])
+handleRet _ op = mkWorld op
+
+getFArgs : NamedCExp -> Core (List (NamedCExp, NamedCExp))
+getFArgs (NmCon fc _ _ (Just 0) _) = pure []
+getFArgs (NmCon fc _ _ (Just 1) [ty, val, rest]) = pure $ (ty, val) :: !(getFArgs rest)
+getFArgs arg = throw (GenericMsg (getFC arg) ("Badly formed c call argument list " ++ show arg))
+
+export
+chezExtPrim : SortedSet Name -> LazyExprProc -> Nat -> ExtPrim -> List NamedCExp -> Core Builder
+chezExtPrim cs schLazy i GetField [NmPrimVal _ (Str s), _, _, struct,
+                                   NmPrimVal _ (Str fld), _]
+    = do structsc <- schExp cs (chezExtPrim cs schLazy) chezString schLazy 0 struct
+         pure $ "(ftype-ref " ++ fromString s ++ " (" ++ fromString fld ++ ") " ++ structsc ++ ")"
+chezExtPrim cs schLazy i GetField [_,_,_,_,_,_]
+    = pure "(blodwen-error-quit \"bad getField\")"
+chezExtPrim cs schLazy i SetField [NmPrimVal _ (Str s), _, _, struct,
+                        NmPrimVal _ (Str fld), _, val, world]
+    = do structsc <- schExp cs (chezExtPrim cs schLazy) chezString schLazy 0 struct
+         valsc <- schExp cs (chezExtPrim cs schLazy) chezString schLazy 0 val
+         pure $ mkWorld $
+            "(ftype-set! " ++ fromString s ++ " (" ++ fromString fld ++ ") " ++ structsc ++
+            " " ++ valsc ++ ")"
+chezExtPrim cs schLazy i SetField [_,_,_,_,_,_,_,_]
+    = pure "(blodwen-error-quit \"bad setField\")"
+chezExtPrim cs schLazy i SysCodegen []
+    = pure $ "\"chez\""
+chezExtPrim cs schLazy i OnCollect [_, p, c, world]
+    = do p' <- schExp cs (chezExtPrim cs schLazy) chezString schLazy 0 p
+         c' <- schExp cs (chezExtPrim cs schLazy) chezString schLazy 0 c
+         pure $ mkWorld $ "(blodwen-register-object " ++ p' ++ " " ++ c' ++ ")"
+chezExtPrim cs schLazy i OnCollectAny [p, c, world]
+    = do p' <- schExp cs (chezExtPrim cs schLazy) chezString schLazy 0 p
+         c' <- schExp cs (chezExtPrim cs schLazy) chezString schLazy 0 c
+         pure $ mkWorld $ "(blodwen-register-object " ++ p' ++ " " ++ c' ++ ")"
+chezExtPrim cs schLazy i prim args
+    = schExtCommon cs (chezExtPrim cs schLazy) chezString schLazy i prim args
 
 -- Reference label for keeping track of loaded external libraries
 export
@@ -502,10 +501,12 @@ compileToSS c prof appdir tm outfile
          fgndefs <- traverse (getFgnCall version) ndefs
          loadlibs <- traverse (locateLib appdir) (mapMaybe fst fgndefs)
 
+         let schLazy = if getWeakMemoLazy ds then weakMemoLaziness else defaultLaziness
+
          (sortedDefs, constants) <- sortDefs ndefs
-         compdefs <- logTime 3 "Print as scheme" $ traverse (getScheme constants (chezExtPrim constants) chezString) sortedDefs
+         compdefs <- logTime 3 "Print as scheme" $ traverse (getScheme constants (chezExtPrim constants schLazy) chezString schLazy) sortedDefs
          let code = concat (map snd fgndefs) ++ concat compdefs
-         main <- schExp constants (chezExtPrim constants) chezString 0 ctm
+         main <- schExp constants (chezExtPrim constants schLazy) chezString schLazy 0 ctm
          support <- readDataFile "chez/support.ss"
          extraRuntime <- getExtraRuntime ds
          let scm = concat $ the (List _)
@@ -552,7 +553,7 @@ compileToSSInc c mods libs appdir tm outfile
          loadlibs <- traverse (map fromString . loadLib appdir) (nub libs)
          loadsos <- traverse (map fromString . loadSO appdir) (nub mods)
 
-         main <- schExp empty (chezExtPrim empty) chezString 0 ctm
+         main <- schExp empty (chezExtPrim empty defaultLaziness) chezString defaultLaziness 0 ctm
          support <- readDataFile "chez/support.ss"
 
          let scm = schHeader chez [] False ++
@@ -692,7 +693,7 @@ incCompile c s sourceFile
                version <- coreLift $ chezVersion chez
                fgndefs <- traverse (getFgnCall version) ndefs
                (sortedDefs, constants) <- sortDefs ndefs
-               compdefs <- traverse (getScheme empty (chezExtPrim empty) chezString) sortedDefs
+               compdefs <- traverse (getScheme empty (chezExtPrim empty defaultLaziness) chezString defaultLaziness) sortedDefs
                let code = concat $ map snd fgndefs ++ compdefs
                Right () <- coreLift $ writeFile ssFile $ build code
                   | Left err => throw (FileErr ssFile err)

src/Compiler/Scheme/ChezSep.idr

@@ -217,7 +217,7 @@ compileToSS c chez appdir tm = do
       let footer = ")"
 
       fgndefs <- traverse (Chez.getFgnCall version) cu.definitions
-      compdefs <- traverse (getScheme empty (Chez.chezExtPrim empty) Chez.chezString) cu.definitions
+      compdefs <- traverse (getScheme empty (Chez.chezExtPrim empty defaultLaziness) Chez.chezString defaultLaziness) cu.definitions
       loadlibs <- traverse (map fromString . loadLib appdir) (mapMaybe fst fgndefs)
 
       -- write the files
@@ -234,7 +234,7 @@ compileToSS c chez appdir tm = do
     pure (MkChezLib chezLib hashChanged)
 
   -- main module
-  main <- schExp empty (Chez.chezExtPrim empty) Chez.chezString 0 ctm
+  main <- schExp empty (Chez.chezExtPrim empty defaultLaziness) Chez.chezString defaultLaziness 0 ctm
   Core.writeFile (appdir </> "mainprog.ss") $ build $ sepBy "\n"
     [ schHeader (map snd libs) [lib.name | lib <- chezLibs]
     , collectRequestHandler

src/Compiler/Scheme/Common.idr

@@ -317,9 +317,28 @@ var _ = False
 getScrutineeTemp : Nat -> Builder
 getScrutineeTemp i = fromString $ "sc" ++ show i
 
+public export
+record LazyExprProc where
+  constructor MkLazyExprProc
+  processDelay : Builder -> Builder
+  processForce : Builder -> Builder
+
+public export
+defaultLaziness : LazyExprProc
+defaultLaziness = MkLazyExprProc
+  (\expr => "(lambda () " ++ expr ++ ")")
+  (\expr => "(" ++ expr ++ ")")
+
+public export
+weakMemoLaziness : LazyExprProc
+weakMemoLaziness = MkLazyExprProc
+  (\expr => "(blodwen-delay-lazy (lambda () " ++ expr ++ "))")
+  (\expr => "(blodwen-force-lazy " ++ expr ++ ")")
+
 parameters (constants : SortedSet Name,
             schExtPrim : Nat -> ExtPrim -> List NamedCExp -> Core Builder,
-            schString : String -> Builder)
+            schString : String -> Builder,
+            schLazy : LazyExprProc)
   showTag : Name -> Maybe Int -> Builder
   showTag n (Just i) = showB i
   showTag n Nothing = schString (show n)
@@ -562,8 +581,8 @@ parameters (constants : SortedSet Name,
         = schExtPrim i (toPrim p) args
     schExp i (NmForce _ _ (NmApp fc x@(NmRef _ _) []))
        = pure $ "(force " ++ !(schExp i x) ++ ")" -- Special version for memoized toplevel lazy definitions
-    schExp i (NmForce fc lr t) = pure $ "(" ++ !(schExp i t) ++ ")"
-    schExp i (NmDelay fc lr t) = pure $ "(lambda () " ++ !(schExp i t) ++ ")"
+    schExp i (NmForce fc lr t) = pure $ schLazy.processForce !(schExp i t)
+    schExp i (NmDelay fc lr t) = pure $ schLazy.processDelay !(schExp i t)
     schExp i (NmConCase fc sc alts def)
         = cond [(recordCase alts, schRecordCase i sc alts def),
                 (maybeCase alts, schMaybeCase i sc alts def),
@@ -678,6 +697,7 @@ getScheme : {auto c : Ref Ctxt Defs} ->
             (constants  : SortedSet Name) ->
             (schExtPrim : Nat -> ExtPrim -> List NamedCExp -> Core Builder) ->
             (schString : String -> Builder) ->
+            (schLazy : LazyExprProc) ->
             (Name, FC, NamedDef) -> Core Builder
-getScheme constants schExtPrim schString (n, fc, d)
-    = schDef constants schExtPrim schString n d
+getScheme constants schExtPrim schString schLazy (n, fc, d)
+    = schDef constants schExtPrim schString schLazy n d

src/Compiler/Scheme/Gambit.idr

@@ -74,35 +74,34 @@ showGambitString (c :: cs) acc = showGambitChar c $ showGambitString cs acc
 gambitString : String -> Builder
 gambitString cs = "\"" ++ showGambitString (unpack cs) "\""
 
-mutual
-  handleRet : CFType -> Builder -> Builder
-  handleRet CFUnit op = op ++ " " ++ mkWorld (schConstructor gambitString (UN $ Basic "") (Just 0) [])
-  handleRet _ op = mkWorld op
-
-  getFArgs : NamedCExp -> Core (List (NamedCExp, NamedCExp))
-  getFArgs (NmCon fc _ _ (Just 0) _) = pure []
-  getFArgs (NmCon fc _ _ (Just 1) [ty, val, rest]) = pure $ (ty, val) :: !(getFArgs rest)
-  getFArgs arg = throw (GenericMsg (getFC arg) ("Badly formed c call argument list " ++ show arg))
-
-  gambitPrim : SortedSet Name -> Nat -> ExtPrim -> List NamedCExp -> Core Builder
-  gambitPrim cs i GetField [NmPrimVal _ (Str s), _, _, struct,
-                         NmPrimVal _ (Str fld), _]
-      = do structsc <- schExp cs (gambitPrim cs) gambitString 0 struct
-           pure $ "(" ++ fromString s ++ "-" ++ fromString fld ++ " " ++ structsc ++ ")"
-  gambitPrim cs i GetField [_,_,_,_,_,_]
-      = pure "(error \"bad getField\")"
-  gambitPrim cs i SetField [NmPrimVal _ (Str s), _, _, struct,
-                         NmPrimVal _ (Str fld), _, val, world]
-      = do structsc <- schExp cs (gambitPrim cs) gambitString 0 struct
-           valsc <- schExp cs (gambitPrim cs) gambitString 0 val
-           pure $ mkWorld $
-                "(" ++ fromString s ++ "-" ++ fromString fld ++ "-set! " ++ structsc ++ " " ++ valsc ++ ")"
-  gambitPrim cs i SetField [_,_,_,_,_,_,_,_]
-      = pure "(error \"bad setField\")"
-  gambitPrim cs i SysCodegen []
-      = pure $ "\"gambit\""
-  gambitPrim cs i prim args
-      = schExtCommon cs (gambitPrim cs) gambitString i prim args
+handleRet : CFType -> Builder -> Builder
+handleRet CFUnit op = op ++ " " ++ mkWorld (schConstructor gambitString (UN $ Basic "") (Just 0) [])
+handleRet _ op = mkWorld op
+
+getFArgs : NamedCExp -> Core (List (NamedCExp, NamedCExp))
+getFArgs (NmCon fc _ _ (Just 0) _) = pure []
+getFArgs (NmCon fc _ _ (Just 1) [ty, val, rest]) = pure $ (ty, val) :: !(getFArgs rest)
+getFArgs arg = throw (GenericMsg (getFC arg) ("Badly formed c call argument list " ++ show arg))
+
+gambitPrim : SortedSet Name -> LazyExprProc -> Nat -> ExtPrim -> List NamedCExp -> Core Builder
+gambitPrim cs schLazy i GetField [NmPrimVal _ (Str s), _, _, struct,
+                       NmPrimVal _ (Str fld), _]
+    = do structsc <- schExp cs (gambitPrim cs schLazy) gambitString schLazy 0 struct
+         pure $ "(" ++ fromString s ++ "-" ++ fromString fld ++ " " ++ structsc ++ ")"
+gambitPrim cs schLazy i GetField [_,_,_,_,_,_]
+    = pure "(error \"bad getField\")"
+gambitPrim cs schLazy i SetField [NmPrimVal _ (Str s), _, _, struct,
+                       NmPrimVal _ (Str fld), _, val, world]
+    = do structsc <- schExp cs (gambitPrim cs schLazy) gambitString schLazy 0 struct
+         valsc <- schExp cs (gambitPrim cs schLazy) gambitString schLazy 0 val
+         pure $ mkWorld $
+              "(" ++ fromString s ++ "-" ++ fromString fld ++ "-set! " ++ structsc ++ " " ++ valsc ++ ")"
+gambitPrim cs schLazy i SetField [_,_,_,_,_,_,_,_]
+    = pure "(error \"bad setField\")"
+gambitPrim cs schLazy i SysCodegen []
+    = pure $ "\"gambit\""
+gambitPrim cs schLazy i prim args
+    = schExtCommon cs (gambitPrim cs schLazy) gambitString schLazy i prim args
 
 -- Reference label for keeping track of loaded external libraries
 data Loaded : Type where
@@ -365,16 +364,18 @@ compileToSCM c tm outfile
          -- let tags = nameTags cdata
          let ctm = forget (mainExpr cdata)
 
+         ds <- getDirectives Gambit
+         let schLazy = if getWeakMemoLazy ds then weakMemoLaziness else defaultLaziness
+
          defs <- get Ctxt
          l <- newRef {t = List String} Loaded []
          s <- newRef {t = List String} Structs []
          fgndefs <- traverse getFgnCall ndefs
          (sortedDefs, constants) <- sortDefs ndefs
-         compdefs <- traverse (getScheme constants (gambitPrim constants) gambitString) ndefs
+         compdefs <- traverse (getScheme constants (gambitPrim constants schLazy) gambitString schLazy) ndefs
          let code = concat (map snd fgndefs) ++ concat compdefs
-         main <- schExp constants (gambitPrim constants) gambitString 0 ctm
+         main <- schExp constants (gambitPrim constants schLazy) gambitString schLazy 0 ctm
          support <- readDataFile "gambit/support.scm"
-         ds <- getDirectives Gambit
          extraRuntime <- getExtraRuntime ds
          foreign <- readDataFile "gambit/foreign.scm"
          let scm = sepBy "\n" [schHeader, fromString support, fromString extraRuntime, fromString foreign, code, main]