Fast nockma eval (#2580)

Adds annotations to cells to indicate that it is a call to the stdlib
and might be evaluated faster in the Haskell evaluator.

The syntax for stdlib calls is as follows:
```
[stdlib@add args@<args-term> <left-term> <right-term>]
```
where `add` is the name of the function being called, `<args-term>` is a
nockma term that points to the position of the arguments, and
`<left-term>` and `<right-term>` are the actual components of the cell.

.github/workflows/ci.yml

@@ -45,9 +45,13 @@ jobs:
         with:
           version: 0.5.3.0
           extra-args: >-
-            --ghc-opt -XDerivingStrategies --ghc-opt -XImportQualifiedPost
-            --ghc-opt -XMultiParamTypeClasses --ghc-opt -XStandaloneDeriving
-            --ghc-opt -XTemplateHaskell --ghc-opt -XUnicodeSyntax
+            --ghc-opt -XDerivingStrategies
+            --ghc-opt -XImportQualifiedPost
+            --ghc-opt -XMultiParamTypeClasses
+            --ghc-opt -XPatternSynonyms
+            --ghc-opt -XStandaloneDeriving
+            --ghc-opt -XTemplateHaskell
+            --ghc-opt -XUnicodeSyntax
 
   build-and-test-linux:
     runs-on: ubuntu-22.04

Makefile

@@ -93,6 +93,7 @@ ormolu:
 		--ghc-opt -XStandaloneDeriving \
 		--ghc-opt -XUnicodeSyntax \
 		--ghc-opt -XDerivingStrategies \
+		--ghc-opt -XPatternSynonyms \
 		--ghc-opt -XMultiParamTypeClasses  \
 		--ghc-opt -XTemplateHaskell \
 		--ghc-opt -XImportQualifiedPost \

app/Commands/Dev/Nockma/Eval.hs

@@ -2,6 +2,7 @@ module Commands.Dev.Nockma.Eval where
 
 import Commands.Base hiding (Atom)
 import Commands.Dev.Nockma.Eval.Options
+import Juvix.Compiler.Nockma.Evaluator.Options
 import Juvix.Compiler.Nockma.Pretty
 import Juvix.Compiler.Nockma.Translation.FromAsm
 import Juvix.Compiler.Nockma.Translation.FromSource qualified as Nockma
@@ -13,7 +14,10 @@ runCommand opts = do
   case parsedTerm of
     Left err -> exitJuvixError (JuvixError err)
     Right (TermCell c) -> do
-      res <- runOutputSem @(Term Natural) (say . ppTrace) (evalCompiledNock' (c ^. cellLeft) (c ^. cellRight))
+      res <-
+        runReader defaultEvalOptions
+          . runOutputSem @(Term Natural) (say . ppTrace)
+          $ evalCompiledNock' (c ^. cellLeft) (c ^. cellRight)
       ret <- getReturn res
       putStrLn (ppPrint ret)
     Right TermAtom {} -> exitFailMsg "Expected nockma input to be a cell"

app/Commands/Dev/Nockma/Repl.hs

@@ -8,6 +8,7 @@ import Control.Exception (throwIO)
 import Control.Monad.State.Strict qualified as State
 import Data.String.Interpolate (__i)
 import Juvix.Compiler.Nockma.Evaluator (NockEvalError, evalRepl, fromReplTerm, programAssignments)
+import Juvix.Compiler.Nockma.Evaluator.Options
 import Juvix.Compiler.Nockma.Language
 import Juvix.Compiler.Nockma.Pretty (ppPrint)
 import Juvix.Compiler.Nockma.Pretty qualified as Nockma
@@ -133,9 +134,10 @@ evalStatement = \case
     prog <- getProgram
     et <-
       liftIO
-        $ runM
-          . runError @(ErrNockNatural Natural)
-          . runError @NockEvalError
+        . runM
+        . runReader defaultEvalOptions
+        . runError @(ErrNockNatural Natural)
+        . runError @NockEvalError
         $ evalRepl (putStrLn . Nockma.ppTrace) prog s t
     case et of
       Left e -> error (show e)

package.yaml

@@ -143,6 +143,7 @@ default-extensions:
   - NoFieldSelectors
   - NoImplicitPrelude
   - OverloadedStrings
+  - PatternSynonyms
   - QuasiQuotes
   - RecordWildCards
   - TemplateHaskell

src/Juvix/Compiler/Nockma/Evaluator.hs

@@ -1,10 +1,12 @@
 module Juvix.Compiler.Nockma.Evaluator
   ( module Juvix.Compiler.Nockma.Evaluator,
     module Juvix.Compiler.Nockma.Evaluator.Error,
+    module Juvix.Compiler.Nockma.Evaluator.Options,
   )
 where
 
 import Juvix.Compiler.Nockma.Evaluator.Error
+import Juvix.Compiler.Nockma.Evaluator.Options
 import Juvix.Compiler.Nockma.Language
 import Juvix.Compiler.Nockma.Pretty
 import Juvix.Prelude hiding (Atom, Path)
@@ -45,7 +47,7 @@ subTermT = go
 subTerm :: (Member (Error NockEvalError) r) => Term a -> Path -> Sem r (Term a)
 subTerm term pos = do
   case term ^? subTermT pos of
-    Nothing -> throw @NockEvalError (error "")
+    Nothing -> throw (InvalidPath "subterm")
     Just t -> return t
 
 setSubTerm :: (Member (Error NockEvalError) r) => Term a -> Path -> Term a -> Sem r (Term a)
@@ -55,11 +57,25 @@ setSubTerm term pos repTerm =
           | isNothing (getFirst old) -> throw @NockEvalError (error "")
           | otherwise -> return new
 
-parseCell :: forall r a. (Members '[Error NockEvalError, Error (ErrNockNatural a)] r, NockNatural a) => Cell a -> Sem r (ParsedCell a)
+parseCell ::
+  forall r a.
+  (Members '[Error NockEvalError, Error (ErrNockNatural a)] r, NockNatural a) =>
+  Cell a ->
+  Sem r (ParsedCell a)
 parseCell c = case c ^. cellLeft of
-  TermAtom a -> ParsedOperatorCell <$> parseOperatorCell a (c ^. cellRight)
+  TermAtom a -> operatorOrStdlibCall a (c ^. cellRight) (c ^. cellInfo . unIrrelevant)
   TermCell l -> return (ParsedAutoConsCell (AutoConsCell l (c ^. cellRight)))
   where
+    operatorOrStdlibCall :: Atom a -> Term a -> Maybe (StdlibCall a) -> Sem r (ParsedCell a)
+    operatorOrStdlibCall a t mcall = do
+      opCell <- parseOperatorCell a t
+      return $ case mcall of
+        Nothing -> ParsedOperatorCell opCell
+        Just call -> ParsedStdlibCallCell (parseStdlibCall opCell call)
+
+    parseStdlibCall :: OperatorCell a -> StdlibCall a -> StdlibCallCell a
+    parseStdlibCall op call = StdlibCallCell call op
+
     parseOperatorCell :: Atom a -> Term a -> Sem r (OperatorCell a)
     parseOperatorCell a t = do
       op <- nockOp a
@@ -84,7 +100,7 @@ programAssignments mprog =
 -- | The stack provided in the replExpression has priority
 evalRepl ::
   forall r a.
-  (PrettyCode a, Members '[Error NockEvalError, Error (ErrNockNatural a)] r, NockNatural a) =>
+  (PrettyCode a, Integral a, Members '[Reader EvalOptions, Error NockEvalError, Error (ErrNockNatural a)] r, NockNatural a) =>
   (Term a -> Sem r ()) ->
   Maybe (Program a) ->
   Maybe (Term a) ->
@@ -105,20 +121,57 @@ evalRepl handleTrace mprog defaultStack expr = do
     namedTerms :: HashMap Text (Term a)
     namedTerms = programAssignments mprog
 
-eval :: forall r a. (PrettyCode a, Members '[Output (Term a), Error NockEvalError, Error (ErrNockNatural a)] r, NockNatural a) => Term a -> Term a -> Sem r (Term a)
+eval ::
+  forall r a.
+  (PrettyCode a, Integral a, Members '[Reader EvalOptions, Output (Term a), Error NockEvalError, Error (ErrNockNatural a)] r, NockNatural a) =>
+  Term a ->
+  Term a ->
+  Sem r (Term a)
 eval stack = \case
   TermAtom a -> throw (ExpectedCell ("eval " <> ppTrace a))
-  TermCell c -> do
-    pc <- parseCell c
-    case pc of
+  TermCell c ->
+    parseCell c >>= \case
       ParsedAutoConsCell a -> goAutoConsCell a
       ParsedOperatorCell o -> goOperatorCell o
+      ParsedStdlibCallCell o -> do
+        ignore <- asks (^. evalIgnoreStdlibCalls)
+        if
+            | ignore -> goOperatorCell (o ^. stdlibCallRaw)
+            | otherwise -> goStdlibCall (o ^. stdlibCallCell)
   where
+    goStdlibCall :: StdlibCall a -> Sem r (Term a)
+    goStdlibCall StdlibCall {..} = do
+      args' <- eval stack _stdlibCallArgs
+      let binArith :: (a -> a -> a) -> Sem r (Term a)
+          binArith f = case args' of
+            TCell (TAtom l) (TAtom r) -> return (TCell (TAtom (f l r)) stack)
+            _ -> error "expected a cell with two atoms"
+
+          unaArith :: (a -> a) -> Sem r (Term a)
+          unaArith f = case args' of
+            TAtom n -> return (TCell (TAtom (f n)) stack)
+            _ -> error "expected an atom"
+
+          binCmp :: (a -> a -> Bool) -> Sem r (Term a)
+          binCmp f = case args' of
+            TCell (TAtom l) (TAtom r) -> return (TCell (TermAtom (nockBool (f l r))) stack)
+            _ -> error "expected a cell with two atoms"
+
+      case _stdlibCallFunction of
+        StdlibDec -> unaArith pred
+        StdlibAdd -> binArith (+)
+        StdlibMul -> binArith (*)
+        StdlibSub -> binArith (-)
+        StdlibDiv -> binArith div
+        StdlibMod -> binArith mod
+        StdlibLt -> binCmp (<)
+        StdlibLe -> binCmp (<=)
+
     goAutoConsCell :: AutoConsCell a -> Sem r (Term a)
     goAutoConsCell c = do
-      _cellLeft <- eval stack (TermCell (c ^. autoConsCellLeft))
-      _cellRight <- eval stack (c ^. autoConsCellRight)
-      return (TermCell Cell {..})
+      l' <- eval stack (TermCell (c ^. autoConsCellLeft))
+      r' <- eval stack (c ^. autoConsCellRight)
+      return (TermCell (Cell l' r'))
 
     goOperatorCell :: OperatorCell a -> Sem r (Term a)
     goOperatorCell c = case c ^. operatorCellOp of
@@ -149,7 +202,7 @@ eval stack = \case
 
         goOpTrace :: Sem r (Term a)
         goOpTrace = do
-          Cell tr a <- asCell "OpTrace" (c ^. operatorCellTerm)
+          Cell' tr a _ <- asCell "OpTrace" (c ^. operatorCellTerm)
           tr' <- eval stack tr
           output tr'
           eval stack a
@@ -164,13 +217,13 @@ eval stack = \case
         goOpPush = do
           cellTerm <- asCell "OpPush" (c ^. operatorCellTerm)
           l <- eval stack (cellTerm ^. cellLeft)
-          let s = TermCell Cell {_cellLeft = l, _cellRight = stack}
+          let s = TermCell (Cell l stack)
           eval s (cellTerm ^. cellRight)
 
         goOpReplace :: Sem r (Term a)
         goOpReplace = do
-          Cell rot1 t2 <- asCell "OpReplace 1" (c ^. operatorCellTerm)
-          Cell ro t1 <- asCell "OpReplace 2" rot1
+          Cell' rot1 t2 _ <- asCell "OpReplace 1" (c ^. operatorCellTerm)
+          Cell' ro t1 _ <- asCell "OpReplace 2" rot1
           r <- asPath ro
           t1' <- eval stack t1
           t2' <- eval stack t2
@@ -187,7 +240,7 @@ eval stack = \case
         goOpIf = do
           cellTerm <- asCell "OpIf 1" (c ^. operatorCellTerm)
           let t0 = cellTerm ^. cellLeft
-          Cell t1 t2 <- asCell "OpIf 2" (cellTerm ^. cellRight)
+          Cell' t1 t2 _ <- asCell "OpIf 2" (cellTerm ^. cellRight)
           cond <- eval stack t0 >>= asBool
           if
               | cond -> eval stack t1

src/Juvix/Compiler/Nockma/Evaluator/Options.hs

@@ -0,0 +1,12 @@
+module Juvix.Compiler.Nockma.Evaluator.Options where
+
+import Juvix.Prelude.Base
+
+newtype EvalOptions = EvalOptions
+  { _evalIgnoreStdlibCalls :: Bool
+  }
+
+defaultEvalOptions :: EvalOptions
+defaultEvalOptions = EvalOptions False
+
+makeLenses ''EvalOptions