LetRec in Core (#1507)

package.yaml

@@ -4,7 +4,7 @@ license:         GPL-3.0-only
 license-file:    LICENSE
 copyright:       (c) 2022- Heliax AG.
 maintainer:      The PLT Team at Heliax AG <hello@heliax.dev>
-author:          [ Jonathan Prieto-Cubides , Jan Mas Rovira , Paul Cadman , Github's contributors ]
+author:          [ Jonathan Prieto-Cubides , Jan Mas Rovira , Paul Cadman , Lukasz Czajka , Github's contributors ]
 tested-with:     ghc == 9.2.4
 homepage:        https://juvix.org
 bug-reports:     https://github.com/anoma/juvix/issues

src/Juvix/Compiler/Core/Data/InfoTable.hs

@@ -7,7 +7,7 @@ type IdentContext = HashMap Symbol Node
 data InfoTable = InfoTable
   { _identContext :: IdentContext,
     -- `_identMap` is needed only for REPL
-    _identMap :: HashMap Text (Either Symbol Tag),
+    _identMap :: HashMap Text IdentKind,
     _infoMain :: Maybe Symbol,
     _infoIdentifiers :: HashMap Symbol IdentifierInfo,
     _infoInductives :: HashMap Name InductiveInfo,
@@ -27,6 +27,8 @@ emptyInfoTable =
       _infoAxioms = mempty
     }
 
+data IdentKind = IdentSym Symbol | IdentTag Tag
+
 data IdentifierInfo = IdentifierInfo
   { _identifierName :: Name,
     _identifierSymbol :: Symbol,

src/Juvix/Compiler/Core/Data/InfoTableBuilder.hs

@@ -11,18 +11,11 @@ data InfoTableBuilder m a where
   RegisterConstructor :: ConstructorInfo -> InfoTableBuilder m ()
   RegisterIdentNode :: Symbol -> Node -> InfoTableBuilder m ()
   SetIdentArgsInfo :: Symbol -> [ArgumentInfo] -> InfoTableBuilder m ()
-  GetIdent :: Text -> InfoTableBuilder m (Maybe (Either Symbol Tag))
+  GetIdent :: Text -> InfoTableBuilder m (Maybe IdentKind)
   GetInfoTable :: InfoTableBuilder m InfoTable
 
 makeSem ''InfoTableBuilder
 
-hasIdent :: Member InfoTableBuilder r => Text -> Sem r Bool
-hasIdent txt = do
-  i <- getIdent txt
-  case i of
-    Just _ -> return True
-    Nothing -> return False
-
 getConstructorInfo :: Member InfoTableBuilder r => Tag -> Sem r ConstructorInfo
 getConstructorInfo tag = do
   tab <- getInfoTable
@@ -67,10 +60,10 @@ runInfoTableBuilder tab =
         return (UserTag (s ^. stateNextUserTag - 1))
       RegisterIdent ii -> do
         modify' (over stateInfoTable (over infoIdentifiers (HashMap.insert (ii ^. identifierSymbol) ii)))
-        modify' (over stateInfoTable (over identMap (HashMap.insert (ii ^. (identifierName . nameText)) (Left (ii ^. identifierSymbol)))))
+        modify' (over stateInfoTable (over identMap (HashMap.insert (ii ^. (identifierName . nameText)) (IdentSym (ii ^. identifierSymbol)))))
       RegisterConstructor ci -> do
         modify' (over stateInfoTable (over infoConstructors (HashMap.insert (ci ^. constructorTag) ci)))
-        modify' (over stateInfoTable (over identMap (HashMap.insert (ci ^. (constructorName . nameText)) (Right (ci ^. constructorTag)))))
+        modify' (over stateInfoTable (over identMap (HashMap.insert (ci ^. (constructorName . nameText)) (IdentTag (ci ^. constructorTag)))))
       RegisterIdentNode sym node ->
         modify' (over stateInfoTable (over identContext (HashMap.insert sym node)))
       SetIdentArgsInfo sym argsInfo -> do

src/Juvix/Compiler/Core/Evaluator.hs

@@ -9,9 +9,10 @@ module Juvix.Compiler.Core.Evaluator where
 import Control.Exception qualified as Exception
 import Data.HashMap.Strict qualified as HashMap
 import Debug.Trace qualified as Debug
+import GHC.Conc qualified as GHC
 import GHC.Show as S
 import Juvix.Compiler.Core.Data.InfoTable
-import Juvix.Compiler.Core.Error
+import Juvix.Compiler.Core.Error (CoreError (..))
 import Juvix.Compiler.Core.Extra
 import Juvix.Compiler.Core.Info qualified as Info
 import Juvix.Compiler.Core.Info.NoDisplayInfo
@@ -50,7 +51,7 @@ eval :: IdentContext -> Env -> Node -> Node
 eval !ctx !env0 = convertRuntimeNodes . eval' env0
   where
     evalError :: Text -> Node -> a
-    evalError !msg !node = Exception.throw (EvalError msg (Just node))
+    evalError msg node = Exception.throw (EvalError msg (Just node))
 
     eval' :: Env -> Node -> Node
     eval' !env !n = case n of
@@ -65,6 +66,10 @@ eval !ctx !env0 = convertRuntimeNodes . eval' env0
       NCtr (Constr i tag args) -> mkConstr i tag (map' (eval' env) args)
       NLam l@Lambda {} -> Closure env l
       NLet (Let _ v b) -> let !v' = eval' env v in eval' (v' : env) b
+      NRec (LetRec _ vs b) ->
+        let !vs' = map (eval' env') (toList vs)
+            !env' = revAppend vs' env
+         in foldr GHC.pseq (eval' env' b) vs'
       NCase (Case i v bs def) ->
         case eval' env v of
           NCtr (Constr _ tag args) -> branch n env args tag def bs
@@ -77,7 +82,10 @@ eval !ctx !env0 = convertRuntimeNodes . eval' env0
 
     branch :: Node -> Env -> [Node] -> Tag -> Maybe Node -> [CaseBranch] -> Node
     branch n !env !args !tag !def = \case
-      (CaseBranch tag' _ b) : _ | tag' == tag -> eval' (revAppend args env) b
+      (CaseBranch tag' _ b) : _
+        | tag' == tag ->
+            let !env' = revAppend args env
+             in eval' env' b
       _ : bs' -> branch n env args tag def bs'
       [] -> case def of
         Just b -> eval' env b

src/Juvix/Compiler/Core/Extra/Base.hs

@@ -2,6 +2,7 @@ module Juvix.Compiler.Core.Extra.Base where
 
 import Data.Functor.Identity
 import Data.List qualified as List
+import Data.List.NonEmpty (fromList)
 import Juvix.Compiler.Core.Info qualified as Info
 import Juvix.Compiler.Core.Info.BinderInfo
 import Juvix.Compiler.Core.Language
@@ -57,6 +58,12 @@ mkLet i v b = NLet (Let i v b)
 mkLet' :: Node -> Node -> Node
 mkLet' = mkLet Info.empty
 
+mkLetRec :: Info -> NonEmpty Node -> Node -> Node
+mkLetRec i vs b = NRec (LetRec i vs b)
+
+mkLetRec' :: NonEmpty Node -> Node -> Node
+mkLetRec' = mkLetRec Info.empty
+
 mkCase :: Info -> Node -> [CaseBranch] -> Maybe Node -> Node
 mkCase i v bs def = NCase (Case i v bs def)
 
@@ -178,6 +185,14 @@ destruct = \case
   NCtr (Constr i tag args) -> NodeDetails i args (map (const 0) args) (map (const []) args) (`mkConstr` tag)
   NLam (Lambda i b) -> NodeDetails i [b] [1] [fetchBinderInfo i] (\i' args' -> mkLambda i' (hd args'))
   NLet (Let i v b) -> NodeDetails i [v, b] [0, 1] [[], fetchBinderInfo i] (\i' args' -> mkLet i' (hd args') (args' !! 1))
+  NRec (LetRec i vs b) ->
+    let n = length vs
+     in NodeDetails
+          i
+          (b : toList vs)
+          (replicate (n + 1) n)
+          (replicate (n + 1) (getInfoBinders n i))
+          (\i' args' -> mkLetRec i' (fromList (tl args')) (hd args'))
   NCase (Case i v bs Nothing) ->
     let branchBinderNums = map (\(CaseBranch _ k _) -> k) bs
      in NodeDetails

src/Juvix/Compiler/Core/Info/BinderInfo.hs

@@ -3,12 +3,21 @@ module Juvix.Compiler.Core.Info.BinderInfo where
 import Juvix.Compiler.Core.Info qualified as Info
 import Juvix.Compiler.Core.Language
 
+-- | Info about a single binder. Associated with Lambda and Pi.
 newtype BinderInfo = BinderInfo {_infoBinder :: Info}
 
 instance IsInfo BinderInfo
 
-kBinderInfoData :: Key BinderInfo
-kBinderInfoData = Proxy
+kBinderInfo :: Key BinderInfo
+kBinderInfo = Proxy
+
+-- | Info about multiple binders. Associated with LetRec.
+newtype BindersInfo = BindersInfo {_infoBinders :: [Info]}
+
+instance IsInfo BindersInfo
+
+kBindersInfo :: Key BindersInfo
+kBindersInfo = Proxy
 
 newtype CaseBinderInfo = CaseBinderInfo
   { _infoBranchBinders :: [[Info]]
@@ -24,6 +33,12 @@ makeLenses ''CaseBinderInfo
 
 getInfoBinder :: Info -> Info
 getInfoBinder i =
-  case Info.lookup kBinderInfoData i of
+  case Info.lookup kBinderInfo i of
     Just (BinderInfo {..}) -> _infoBinder
     Nothing -> Info.empty
+
+getInfoBinders :: Int -> Info -> [Info]
+getInfoBinders n i =
+  case Info.lookup kBindersInfo i of
+    Just (BindersInfo {..}) -> _infoBinders
+    Nothing -> replicate n Info.empty

src/Juvix/Compiler/Core/Info/NameInfo.hs

@@ -17,3 +17,6 @@ getInfoName i =
   case Info.lookup kNameInfo i of
     Just (NameInfo {..}) -> Just _infoName
     Nothing -> Nothing
+
+setInfoName :: Name -> Info -> Info
+setInfoName = Info.insert . NameInfo

src/Juvix/Compiler/Core/Info/TypeInfo.hs

@@ -17,3 +17,6 @@ getInfoType i =
   case Info.lookup kTypeInfo i of
     Just (TypeInfo {..}) -> Just _infoType
     Nothing -> Nothing
+
+setInfoType :: Type -> Info -> Info
+setInfoType = Info.insert . TypeInfo

src/Juvix/Compiler/Core/Language.hs

@@ -56,6 +56,15 @@ data Lambda = Lambda {_lambdaInfo :: !Info, _lambdaBody :: !Node}
 -- purposes of ML-polymorphic / dependent type checking or code generation!
 data Let = Let {_letInfo :: !Info, _letValue :: !Node, _letBody :: !Node}
 
+-- | Represents a block of mutually recursive local definitions. Both in the
+-- body and in the values `length _letRecValues` implicit binders are introduced
+-- which hold the functions/values being defined.
+data LetRec = LetRec
+  { _letRecInfo :: !Info,
+    _letRecValues :: !(NonEmpty Node),
+    _letRecBody :: !Node
+  }
+
 -- | One-level case matching on the tag of a data constructor: `Case value
 -- branches default`. `Case` is lazy: only the selected branch is evaluated.
 data Case = Case
@@ -100,12 +109,13 @@ data Node
   | NCtr {-# UNPACK #-} !Constr
   | NLam {-# UNPACK #-} !Lambda
   | NLet {-# UNPACK #-} !Let
+  | NRec {-# UNPACK #-} !LetRec
   | NCase {-# UNPACK #-} !Case
   | NPi {-# UNPACK #-} !Pi
   | NUniv {-# UNPACK #-} !Univ
   | NTyp {-# UNPACK #-} !TypeConstr
   | NDyn !Dynamic -- Dynamic is already a newtype, so it's unpacked.
-  | -- Evaluation only: `Closure env body`
+  | -- Evaluation only: `Closure env body`.
     Closure
       { _closureEnv :: !Env,
         _closureLambda :: {-# UNPACK #-} !Lambda
@@ -174,6 +184,9 @@ instance HasAtomicity Lambda where
 instance HasAtomicity Let where
   atomicity _ = Aggregate lambdaFixity
 
+instance HasAtomicity LetRec where
+  atomicity _ = Aggregate lambdaFixity
+
 instance HasAtomicity Case where
   atomicity _ = Aggregate lambdaFixity
 
@@ -199,6 +212,7 @@ instance HasAtomicity Node where
     NCtr x -> atomicity x
     NLam x -> atomicity x
     NLet x -> atomicity x
+    NRec x -> atomicity x
     NCase x -> atomicity x
     NPi x -> atomicity x
     NUniv x -> atomicity x
@@ -233,6 +247,9 @@ instance Eq Lambda where
 instance Eq Let where
   (Let _ v1 b1) == (Let _ v2 b2) = v1 == v2 && b1 == b2
 
+instance Eq LetRec where
+  (LetRec _ vs1 b1) == (LetRec _ vs2 b2) = vs1 == vs2 && b1 == b2
+
 instance Eq Case where
   (Case _ v1 bs1 def1) == (Case _ v2 bs2 def2) = v1 == v2 && bs1 == bs2 && def1 == def2
 
@@ -255,6 +272,7 @@ makeLenses ''App
 makeLenses ''BuiltinApp
 makeLenses ''Constr
 makeLenses ''Let
+makeLenses ''LetRec
 makeLenses ''Case
 makeLenses ''Pi
 makeLenses ''Univ

src/Juvix/Compiler/Core/Pretty/Base.hs

@@ -5,6 +5,7 @@ module Juvix.Compiler.Core.Pretty.Base
   )
 where
 
+import Data.List qualified as List
 import Juvix.Compiler.Core.Extra
 import Juvix.Compiler.Core.Info qualified as Info
 import Juvix.Compiler.Core.Info.BinderInfo as BinderInfo
@@ -63,6 +64,7 @@ instance PrettyCode Tag where
     UserTag tag -> return $ kwUnnamedConstr <> pretty tag
 
 instance PrettyCode Node where
+  ppCode :: forall r. Member (Reader Options) r => Node -> Sem r (Doc Ann)
   ppCode node = case node of
     NVar Var {..} ->
       case Info.lookup kNameInfo _varInfo of
@@ -102,7 +104,7 @@ instance PrettyCode Node where
       b <- ppCode body
       return $ foldl' (flip (<+>)) b pplams
       where
-        ppLam :: Member (Reader Options) r => Info -> Sem r (Doc Ann)
+        ppLam :: Info -> Sem r (Doc Ann)
         ppLam i =
           case getInfoName (getInfoBinder i) of
             Just name -> do
@@ -117,6 +119,28 @@ instance PrettyCode Node where
       v' <- ppCode _letValue
       b' <- ppCode _letBody
       return $ kwLet <+> n' <+> kwAssign <+> v' <+> kwIn <+> b'
+    NRec LetRec {..} -> do
+      let n = length _letRecValues
+      ns <- mapM getName (getInfoBinders n _letRecInfo)
+      vs <- mapM ppCode _letRecValues
+      b' <- ppCode _letRecBody
+      if
+          | length ns == 1 ->
+              return $ kwLetRec <+> List.head ns <+> kwAssign <+> head vs <+> kwIn <+> b'
+          | otherwise ->
+              let bss =
+                    indent' $
+                      align $
+                        concatWith (\a b -> a <> kwSemicolon <> line <> b) $
+                          zipWithExact (\name val -> name <+> kwAssign <+> val) ns (toList vs)
+                  nss = enclose kwSquareL kwSquareR (concatWith (<+>) ns)
+               in return $ kwLetRec <> nss <> line <> bss <> line <> kwIn <> line <> b'
+      where
+        getName :: Info -> Sem r (Doc Ann)
+        getName i =
+          case getInfoName i of
+            Just name -> ppCode name
+            Nothing -> return kwQuestion
     NCase Case {..} -> do
       bns <-
         case Info.lookup kCaseBinderInfo _caseInfo of
@@ -204,6 +228,12 @@ ppLRExpression associates fixlr e =
 {--------------------------------------------------------------------------------}
 {- keywords -}
 
+kwSquareL :: Doc Ann
+kwSquareL = delimiter "["
+
+kwSquareR :: Doc Ann
+kwSquareR = delimiter "]"
+
 kwDeBruijnVar :: Doc Ann
 kwDeBruijnVar = keyword Str.deBruijnVar
 
@@ -237,6 +267,9 @@ kwDiv = keyword Str.div
 kwMod :: Doc Ann
 kwMod = keyword Str.mod
 
+kwLetRec :: Doc Ann
+kwLetRec = keyword Str.letrec_
+
 kwCase :: Doc Ann
 kwCase = keyword Str.case_