Add "let memo" syntax for automatic memoization

7 files changed, 119 insertions, 39 deletions

diff --git a/examples/fibo.milch b/examples/fibo.milch
index c08fdd4..f6a14a8 100644
--- a/examples/fibo.milch
+++ b/examples/fibo.milch
@@ -1,7 +1,7 @@
-(let fibo (\[n]
+(let memo fibo (\[n]
     (match n
         0  0
         1  1
         _  (+ (fibo (- n 1)) (fibo (- n 2))))))
 
-(fibo 30)
+(fibo 50)
diff --git a/src/Builtins.hs b/src/Builtins.hs
index 3b92926..079195f 100644
--- a/src/Builtins.hs
+++ b/src/Builtins.hs
@@ -5,11 +5,12 @@ import qualified Data.Map as M
 import qualified Data.Text as T
 import qualified Text.Read as TR
 import qualified Data.List as L
+import qualified Data.Bifunctor as B
 import Control.Monad.Except
 import Utils
 
 builtinEnv :: Env
-builtinEnv = M.fromList [
+builtinEnv = M.fromList $ map (B.second Regular) [
     -- arithmetic
     builtinAdd2,
     builtinSubtract2,
diff --git a/src/Interpreter.hs b/src/Interpreter.hs
index 34fc6bb..585667b 100644
--- a/src/Interpreter.hs
+++ b/src/Interpreter.hs
@@ -177,19 +177,25 @@ evaluateLet asts = do
     d <- getDepth
     when (d > 1) $ throwL (astPos letAst) $ "let can only be called on the top level or in a function definition, current depth: " ++ show d
 
-    (symbol, value) <- case args of
-        [AST { an = ASTSymbol "lazy" }, AST { an = ASTSymbol symbol' }, value'] -> do
-            return (symbol', value')
-        [AST { an = ASTSymbol symbol' }, value'] -> do
+    case args of
+        [AST { an = ASTSymbol "lazy" }, AST { an = ASTSymbol sym }, val] -> do
+            env <- getEnv
+            when (MB.isJust $ resolveSymbol sym env) $ throwL (astPos letAst) $ "symbol already defined: " ++ sym
+            insertEnv sym $ Regular val
+            return $ letAst { an = ASTUnit }
+        [AST { an = ASTSymbol "memo" }, AST { an = ASTSymbol sym }, val] -> do
+            env <- getEnv
+            when (MB.isJust $ resolveSymbol sym env) $ throwL (astPos letAst) $ "symbol already defined: " ++ sym
+            insertEnv sym $ Memoized M.empty val
+            return $ letAst { an = ASTUnit }
+        [AST { an = ASTSymbol sym }, value'] -> do
             evaledValue <- evaluate value'
-            return (symbol', evaledValue)
+            env <- getEnv
+            when (MB.isJust $ resolveSymbol sym env) $ throwL (astPos letAst) $ "symbol already defined: " ++ sym
+            insertEnv sym $ Regular evaledValue
+            return $ letAst { an = ASTUnit }
         other -> throwL (astPos $ head other) $ "let called with invalid args " ++ show other
 
-    env <- getEnv
-    when (MB.isJust $ resolveSymbol symbol env) $ throwL (astPos letAst) $ "symbol already defined: " ++ symbol
-    insertEnv symbol value
-    return $ letAst { an = ASTUnit }
-
 evaluateDebugEnv :: [AST] -> LContext AST
 evaluateDebugEnv asts = do
     let envAst = head asts
@@ -264,14 +270,14 @@ evaluateRecord asts = do
                 makeFnCreate _ _ = error $ "unreachable: makeFnCreate " ++ fnCreateName
 
             let createFn = makeFnCreate fields []
-            insertEnv fnCreateName $ recordAst { an = ASTFunction Pure createFn }
+            insertEnv fnCreateName $ Regular $ recordAst { an = ASTFunction Pure createFn }
 
             let makeGetFns [] = return $ ()
                 makeGetFns (param:restParams) = do
                     let fnGetName = ns ++ "/" ++ "get-" ++ param
                     let fn = getFn fnGetName
                     let fnAST = makeNonsenseAST $ ASTFunction Pure $ fn
-                    insertEnv fnGetName fnAST
+                    insertEnv fnGetName $ Regular fnAST
                     makeGetFns restParams
 
             makeGetFns fields
@@ -281,7 +287,7 @@ evaluateRecord asts = do
                     let fnSetName = ns ++ "/" ++ "set-" ++ param
                     let fn = setFn fnSetName
                     let fnAST = makeNonsenseAST $ ASTFunction Pure $ fn
-                    insertEnv fnSetName fnAST
+                    insertEnv fnSetName $ Regular fnAST
                     makeSetFns restParams
 
             makeSetFns fields
@@ -315,33 +321,83 @@ evaluateRecord asts = do
                     return $ makeNonsenseAST $ ASTRecord identifier newRecord
                 fn ast2 = throwL (astPos ast2) $ "invalid argument passed to " ++ fnName ++ ": " ++ (show ast2)
 
-evaluateUserFunction :: [AST] -> LContext AST
-evaluateUserFunction children = do
+data ReifyResult
+    = ReifyRegularFunction AST
+    | ReifyMemoizedFunction String AST
+
+reifyFunctionReference :: AST -> LContext ReifyResult
+reifyFunctionReference ref = case ref of
+    AST { an = ASTSymbol sym } -> do
+        env <- getEnv
+        let bindingM = resolveSymbol sym env
+        case bindingM of
+            Just binding -> case binding of
+                Regular bound -> return $ ReifyRegularFunction $ bound
+                Memoized _memoMap bound -> return $ ReifyMemoizedFunction sym $ bound
+            Nothing -> throwL (astPos ref) $ "symbol " ++ sym ++ " not defined in environment"
+    other -> do
+        ret <- evaluate other
+        return $ ReifyRegularFunction ret
+
+unsafeGetMemoMap :: String -> LContext (M.Map [AST] AST)
+unsafeGetMemoMap sym = do
+    env <- getEnv
+    case ((M.!) env sym) of
+        Memoized memoMap _ -> return memoMap
+        _ -> error $ "unreachable: unsafeGetMemoMap " ++ show env ++ ", " ++ sym
+
+evaluateFunctionCall :: [AST] -> LContext AST
+evaluateFunctionCall children = do
     let fnAst = head children
         args = tail children
-    fnEvaled <- evaluate fnAst
-    AST { an = astFn@(ASTFunction fIsPure _) } <- assertIsASTFunction fnEvaled
 
-    checkPurity fIsPure
-    updatePurity fIsPure
+    reifyRes <- reifyFunctionReference fnAst
+
+    case reifyRes of
+        -- ReifyRegularFunction bound -> trace ("not memoed: " ++ show fnAst) $ do
+        ReifyRegularFunction bound -> do
+            evaledBound <- evaluate bound
+            AST { an = astFn@(ASTFunction fIsPure _) } <- assertIsASTFunction evaledBound
+
+            checkPurity fIsPure
+            updatePurity fIsPure
+
+            evaledArgs <- mapM evaluate args
+            result <- curryCall (reverse evaledArgs) astFn
+            return $ fnAst { an = an result }
+        ReifyMemoizedFunction sym bound -> do
+            evaledArgs <- mapM evaluate args
+
+            memoMap <- unsafeGetMemoMap sym
+
+            case (M.lookup evaledArgs memoMap) of
+                Just hit -> return hit
+                Nothing -> do
+                    evaledBound <- evaluate bound
+                    AST { an = astFn@(ASTFunction fIsPure _) } <- assertIsASTFunction evaledBound
+
+                    checkPurity fIsPure
+                    updatePurity fIsPure
 
-    evaledArgs <- mapM evaluate args
-    doubleEvaledArgs <- mapM evaluate evaledArgs
+                    result <- curryCall (reverse evaledArgs) astFn
 
-    result <- curryCall (reverse doubleEvaledArgs) astFn
+                    possiblyUpdatedMemoMap <- unsafeGetMemoMap sym
+                    let newMemoMap = M.insert evaledArgs result possiblyUpdatedMemoMap
 
-    -- todo: maybe remove double eval here? can't remember why it was added
-    return $ fnAst { an = an result }
+                    insertEnv sym $ Memoized newMemoMap bound
+                    return $ fnAst { an = an result }
 
-resolveSymbol :: String -> Env -> Maybe AST
+resolveSymbol :: String -> Env -> Maybe (Binding AST)
 resolveSymbol = M.lookup
 
 evaluateSymbol :: AST -> LContext AST
 evaluateSymbol ast@AST { an = ASTSymbol sym } = do
     env <- getEnv
-    let val = resolveSymbol sym env
-    case val of
-        Just ast' -> return ast'
+    let bindingM = resolveSymbol sym env
+    case bindingM of
+        Just binding -> return $ case binding of
+            Regular v -> v
+            Memoized _ v -> v
         Nothing -> throwL (astPos ast) $ "symbol " ++ sym ++ " not defined in environment"
 evaluateSymbol ast = throwL (astPos ast) $ "unreachable: evaluateSymbol, ast: " ++ show ast
 
@@ -370,7 +426,7 @@ evaluate ast@AST { an = fnc@(ASTFunctionCall args@(x:_)) } =
                 ASTSymbol "record" ->
                     evaluateRecord args
                 _ ->
-                    evaluateUserFunction args
+                    evaluateFunctionCall args
 
         ret <- task `catchError`
             appendError ("when calling function " ++ show x ++ " at " ++ astPos ast)
diff --git a/src/Utils.hs b/src/Utils.hs
index 127e7c0..e4b2333 100644
--- a/src/Utils.hs
+++ b/src/Utils.hs
@@ -29,8 +29,12 @@ data Config = Config {
     configPrintCallStack :: Bool,
     configUseREPL :: Bool
 }
+data Binding a
+    = Regular a
+    | Memoized (M.Map [a] a) a
+    deriving Show
 
-type Env = M.Map String AST
+type Env = M.Map String (Binding AST)
 type Scope = [(String, AST)]
 
 data LState = LState {
@@ -61,7 +65,7 @@ putEnv :: Env -> LContext ()
 putEnv env = do
     modify (\s -> s { stateEnv = env })
 
-insertEnv :: String -> AST -> LContext ()
+insertEnv :: String -> Binding AST -> LContext ()
 insertEnv k v = do
     env <- getEnv
     putEnv $ M.insert k v env
diff --git a/test/Spec.hs b/test/Spec.hs
index 82b0813..e27159c 100644
--- a/test/Spec.hs
+++ b/test/Spec.hs
@@ -32,7 +32,7 @@ parseTests = testGroup "parse" [
     ]
 
 evaluateTests = testGroup "evaluate" [
-     do let env = M.fromList [builtinAdd2] :: Env
+     do let env = makeEnv [builtinAdd2]
         (gotAST, LState { stateEnv = gotEnv }) <- expectSuccessL env $
             evaluate (astFunctionCall [astSymbol "+", astInteger 1, astInteger 2])
 
@@ -42,7 +42,7 @@ evaluateTests = testGroup "evaluate" [
     ]
 
 e2eTests = testGroup "e2e" [
-     do let env = M.fromList [builtinSubtract2] :: Env
+     do let env = makeEnv [builtinSubtract2]
         let script1 = "(let sub2 (\\[a b] (- a b)))\n(sub2 3 2)"
         (gotASTs, LState { stateEnv = gotEnv }) <- expectSuccessL env $
             runInlineScript "<test>" script1
@@ -68,7 +68,7 @@ e2eTests = testGroup "e2e" [
         let expectedLastAST = ast $ ASTInteger 123
         assertEqual "" expectedLastAST (last gotASTs),
 
-     do let env = M.fromList [builtinSortByFirst] :: Env
+     do let env = makeEnv [builtinSortByFirst]
         let script1 = "(sort-by-first [[2 1] [3 2] [1 3]])"
         (gotASTs, _) <- expectSuccessL env $
             runInlineScript "<test>" script1
@@ -79,12 +79,26 @@ e2eTests = testGroup "e2e" [
                  astVector [astInteger 3, astInteger 2]]
         assertEqual "" expectedLastAST (last gotASTs),
 
-     do let env = M.fromList [builtinAdd2, builtinMultiply2] :: Env
+     do let env = makeEnv [builtinAdd2, builtinMultiply2]
         let script1 = "(let f (\\[x] (let y (+ x 1)) (let z (+ 3 y)) (* 2 z)))\n(f 1)"
         (gotASTs, _) <- expectSuccessL env $
             runInlineScript "<test>" script1
 
         let expectedLastAST = astInteger 10
+        assertEqual "" expectedLastAST (last gotASTs),
+
+     do let env = makeEnv [builtinAdd2, builtinSubtract2, ("_", astHole)]
+        let script1 = "(let memo fibo (\\[n]\
+                      \(match n\
+                      \    0  0\
+                      \    1  1\
+                      \    _  (+ (fibo (- n 1)) (fibo (- n 2))))))\
+                      \ \
+                      \(fibo 50)"
+        (gotASTs, _) <- expectSuccessL env $
+            runInlineScript "<test>" script1
+
+        let expectedLastAST = astInteger 12586269025
         assertEqual "" expectedLastAST (last gotASTs)
     ]
 
diff --git a/test/TestUtils.hs b/test/TestUtils.hs
index 86b3db5..1f0d386 100644
--- a/test/TestUtils.hs
+++ b/test/TestUtils.hs
@@ -1,6 +1,8 @@
 {-# OPTIONS_GHC -Wno-missing-export-lists #-}
 module TestUtils where
 
+import qualified Data.Map as M
+import qualified Data.Bifunctor as B
 import Utils
 
 testConfig :: Config
@@ -35,6 +37,9 @@ expectErrorL env lc =
             Left (LException _ err) -> return err
             Right (val, _) -> error $ "unexpected success: " ++ show val
 
+makeEnv :: [(String, AST)] -> Env
+makeEnv = M.fromList . map (B.second Regular)
+
 ast :: ASTNode -> AST
 ast node = AST { an = node }
 
diff --git a/todo.md b/todo.md
index 534abcc..4efb8d4 100644
--- a/todo.md
+++ b/todo.md
@@ -2,6 +2,6 @@
 
 In order of priority
 
-- Memoize pure top-level let-bindings with `let memo`
 - Add a `catch` builtin for catching fatal errors
+- Add :keyword data type, use in core functions
 - Write tests!