Effects for finite and infinite supplies (#2780)

This pr refactors the `Input` effect. It is now meant to be used with
finite input lists.

It also introduces the `StreamOf` effect, which is meant to be used for
infinite supplies.

Both have static implementations so they should add negligible overhead
when used.

src/Juvix/Compiler/Concrete/Translation/FromSource.hs

@@ -239,7 +239,7 @@ runMarkdownModuleParser fpath mk =
 
     parseRestBlocks ::
       forall r'.
-      (Members '[ParserResultBuilder, Error ParserError, Input (Maybe MK.JuvixCodeBlock), State MdModuleBuilder] r') =>
+      (Members '[ParserResultBuilder, Error ParserError, Input MK.JuvixCodeBlock, State MdModuleBuilder] r') =>
       Sem r' ()
     parseRestBlocks = whenJustM input $ \x -> do
       stmts <- parseHelper parseTopStatements x

src/Juvix/Compiler/Concrete/Translation/ImportScanner/FlatParse.hs

@@ -31,7 +31,7 @@ scanner :: Path Abs File -> ByteString -> Result e ScanResult
 scanner fp bs = do
   spansToLocs <$> runParser pPreScanResult bs
   where
-    getInterval :: (Members '[Input (Maybe FileLoc)] r) => Sem r Interval
+    getInterval :: (Members '[Input FileLoc] r) => Sem r Interval
     getInterval = do
       _intervalStart <- inputJust
       _intervalEnd <- inputJust

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

@@ -474,10 +474,6 @@ twoManyChildrenI f i is = \case
   (x : y : xs) -> f i is x y xs
   _ -> impossible
 
-{-# INLINE input' #-}
-input' :: (Members '[Input (Maybe a)] r) => Sem r a
-input' = fmap fromJust input
-
 -- | Destruct a node into NodeDetails. This is an ugly internal function used to
 -- implement more high-level accessors and recursors.
 destruct :: Node -> NodeDetails
@@ -587,11 +583,11 @@ destruct = \case
             [ br : reverse (foldl' (\r b -> manyBinders (take (length r) bi) (b ^. binderType) : r) [] bi)
               | (bi, br) <- branchChildren
             ]
-        mkBranch :: Info -> CaseBranch -> Sem '[Input (Maybe Node)] CaseBranch
+        mkBranch :: Info -> CaseBranch -> Sem '[Input Node] CaseBranch
         mkBranch nfo' br = do
-          b' <- input'
+          b' <- inputJust
           let nBinders = br ^. caseBranchBindersNum
-          tys' <- replicateM nBinders input'
+          tys' <- replicateM nBinders inputJust
           return
             br
               { _caseBranchInfo = nfo',
@@ -600,12 +596,12 @@ destruct = \case
               }
         mkBranches :: [Info] -> [Node] -> [CaseBranch]
         mkBranches is' allNodes' =
-          run $
-            runInputList allNodes' $
-              sequence
-                [ mkBranch ci' br
-                  | (ci', br) <- zipExact is' brs
-                ]
+          run
+            . runInputList allNodes'
+            $ sequence
+              [ mkBranch ci' br
+                | (ci', br) <- zipExact is' brs
+              ]
      in case mdef of
           Nothing ->
             NodeDetails
@@ -651,29 +647,29 @@ destruct = \case
               | br <- branches
             ]
         -- sets the infos and the binder types in the patterns
-        setPatternsInfos :: forall r. (Members '[Input (Maybe Info), Input (Maybe Node)] r) => NonEmpty Pattern -> Sem r (NonEmpty Pattern)
+        setPatternsInfos :: forall r. (Members '[Input Info, Input Node] r) => NonEmpty Pattern -> Sem r (NonEmpty Pattern)
         setPatternsInfos = mapM goPattern
           where
             goPattern :: Pattern -> Sem r Pattern
             goPattern = \case
               PatWildcard x -> do
-                i' <- input'
-                ty <- input'
+                i' <- inputJust
+                ty <- inputJust
                 return (PatWildcard (over patternWildcardBinder (set binderType ty) (set patternWildcardInfo i' x)))
               PatConstr x -> do
-                i' <- input'
-                ty <- input'
+                i' <- inputJust
+                ty <- inputJust
                 args' <- mapM goPattern (x ^. patternConstrArgs)
                 return (PatConstr (over patternConstrBinder (set binderType ty) (set patternConstrInfo i' (set patternConstrArgs args' x))))
      in NodeDetails
           { _nodeInfo = i,
             _nodeSubinfos = branchInfos,
             _nodeChildren = allNodes,
             _nodeReassemble = someChildrenI $ \i' is' chs' ->
-              let mkBranch :: MatchBranch -> Sem '[Input (Maybe Node), Input (Maybe Info)] MatchBranch
+              let mkBranch :: MatchBranch -> Sem '[Input Node, Input Info] MatchBranch
                   mkBranch br = do
-                    bi' <- input'
-                    b' <- input'
+                    bi' <- inputJust
+                    b' <- inputJust
                     pats' <- setPatternsInfos (br ^. matchBranchPatterns)
                     return
                       br

src/Juvix/Compiler/Nockma/Translation/FromTree.hs

@@ -100,7 +100,7 @@ data FunctionInfo = FunctionInfo
     _functionInfoName :: Text
   }
 
-data FunctionCtx = FunctionCtx
+newtype FunctionCtx = FunctionCtx
   { _functionCtxArity :: Natural
   }
 
@@ -520,7 +520,7 @@ compile = \case
                 ClosureTotalArgsNum -> Nothing
                 ClosureArgsNum -> Nothing
                 AnomaGetOrder -> Nothing
-          return $ (opCall "callClosure" (closurePath WrapperCode) newSubject)
+          return (opCall "callClosure" (closurePath WrapperCode) newSubject)
 
 isZero :: Term Natural -> Term Natural
 isZero a = OpEq # a # nockNatLiteral 0
@@ -773,11 +773,11 @@ runCompilerWith opts constrs moduleFuns mainFun = makeAnomaFun
             AnomaGetOrder -> nockNilHere
 
     functionInfos :: HashMap FunctionId FunctionInfo
-    functionInfos = hashMap (run (runInputNaturals (toList <$> userFunctions)))
+    functionInfos = hashMap (run (runStreamOfNaturals (toList <$> userFunctions)))
 
-    userFunctions :: (Members '[Input Natural] r) => Sem r (NonEmpty (FunctionId, FunctionInfo))
+    userFunctions :: (Members '[StreamOf Natural] r) => Sem r (NonEmpty (FunctionId, FunctionInfo))
     userFunctions = forM allFuns $ \CompilerFunction {..} -> do
-      i <- input
+      i <- yield
       return
         ( _compilerFunctionId,
           FunctionInfo

src/Juvix/Prelude/Effects.hs

@@ -3,6 +3,7 @@ module Juvix.Prelude.Effects
     module Juvix.Prelude.Effects.Base,
     module Juvix.Prelude.Effects.Accum,
     module Juvix.Prelude.Effects.Input,
+    module Juvix.Prelude.Effects.StreamOf,
     module Juvix.Prelude.Effects.Bracket,
   )
 where
@@ -12,3 +13,4 @@ import Juvix.Prelude.Effects.Base
 import Juvix.Prelude.Effects.Bracket
 import Juvix.Prelude.Effects.Input
 import Juvix.Prelude.Effects.Output
+import Juvix.Prelude.Effects.StreamOf

src/Juvix/Prelude/Effects/Input.hs

@@ -1,37 +1,38 @@
-{-# OPTIONS_GHC -Wno-unused-type-patterns #-}
+module Juvix.Prelude.Effects.Input
+  ( Input,
+    input,
+    inputJust,
+    peekInput,
+    runInputList,
+  )
+where
 
-module Juvix.Prelude.Effects.Input where
-
-import Data.Stream qualified as Stream
 import Juvix.Prelude.Base.Foundation
 import Juvix.Prelude.Effects.Base
-import Juvix.Prelude.Stream
-
--- TODO make static versions. Finite and infinite.
-data Input (i :: GHCType) :: Effect where
-  Input :: Input i m i
-
-makeEffect ''Input
-
-runInputList :: forall i r a. [i] -> Sem (Input (Maybe i) ': r) a -> Sem r a
-runInputList s = reinterpret (evalState s) $ \case
-  Input -> do
-    x <- gets @[i] nonEmpty
-    case x of
-      Nothing -> return Nothing
-      Just (a :| as) -> do
-        put as
-        return (Just a)
-
-runInputStream :: forall i r a. Stream i -> Sem (Input i ': r) a -> Sem r a
-runInputStream s = reinterpret (evalState s) $ \case
-  Input -> do
-    Stream.Cons a as <- get @(Stream i)
-    put as
-    return a
-
-runInputNaturals :: Sem (Input Natural ': r) a -> Sem r a
-runInputNaturals = runInputStream allNaturals
-
-inputJust :: (Members '[Input (Maybe i)] r) => Sem r i
+import Safe
+
+data Input (i :: GHCType) :: Effect
+
+type instance DispatchOf (Input _) = 'Static 'NoSideEffects
+
+newtype instance StaticRep (Input i) = Input
+  { _unInput :: [i]
+  }
+
+input :: (Member (Input i) r) => Sem r (Maybe i)
+input =
+  stateStaticRep $
+    \case
+      Input [] -> (Nothing, Input [])
+      Input (i : is) -> (Just i, Input is)
+
+peekInput :: (Member (Input i) r) => Sem r (Maybe i)
+peekInput = do
+  Input l <- getStaticRep
+  return (headMay l)
+
+runInputList :: [i] -> Sem (Input i ': r) a -> Sem r a
+runInputList = evalStaticRep . Input
+
+inputJust :: (Members '[Input i] r) => Sem r i
 inputJust = fromMaybe (error "inputJust") <$> input

src/Juvix/Prelude/Effects/StreamOf.hs

@@ -0,0 +1,30 @@
+module Juvix.Prelude.Effects.StreamOf
+  ( StreamOf,
+    yield,
+    runStreamOf,
+    runStreamOfNaturals,
+  )
+where
+
+import Data.Stream
+import Juvix.Prelude.Base.Foundation
+import Juvix.Prelude.Effects.Base
+import Juvix.Prelude.Stream
+
+data StreamOf (i :: GHCType) :: Effect
+
+type instance DispatchOf (StreamOf _) = 'Static 'NoSideEffects
+
+newtype instance StaticRep (StreamOf i) = StreamOf
+  { _unStreamOf :: Stream i
+  }
+
+yield :: (Member (StreamOf i) r) => Sem r i
+yield = stateStaticRep $ \case
+  StreamOf (Cons i is) -> (i, StreamOf is)
+
+runStreamOf :: Stream i -> Sem (StreamOf i ': r) a -> Sem r a
+runStreamOf = evalStaticRep . StreamOf
+
+runStreamOfNaturals :: Sem (StreamOf Natural ': r) a -> Sem r a
+runStreamOfNaturals = runStreamOf allNaturals