break up instances for Node

src/Juvix/Compiler/Core/Evaluator.hs

@@ -43,7 +43,7 @@ instance Show EvalError where
 
 instance Exception.Exception EvalError
 
--- `eval ctx env n` evalues a node `n` whose all free variables point into
+-- | `eval ctx env n` evalues a node `n` whose all free variables point into
 -- `env`. All nodes in `ctx` must be closed. All nodes in `env` must be values.
 -- Invariant for values v: eval ctx env v = v
 eval :: IdentContext -> Env -> Node -> Node
@@ -59,11 +59,11 @@ eval !ctx !env0 = convertRuntimeNodes . eval' env0
       NCst {} -> n
       NApp (App i l r) ->
         case eval' env l of
-          Closure _ env' b -> let !v = eval' env r in eval' (v : env') b
+          Closure env' (Lambda _ b) -> let !v = eval' env r in eval' (v : env') b
           v -> evalError "invalid application" (mkApp i v (substEnv env r))
       NBlt (BuiltinApp _ op args) -> applyBuiltin n env op args
       NCtr (Constr i tag args) -> mkConstr i tag (map (eval' env) args)
-      NLam (Lambda i b) -> Closure i env b
+      NLam l@Lambda {} -> Closure env l
       NLet (Let _ v b) -> let !v' = eval' env v in eval' (v' : env) b
       NCase (Case i v bs def) ->
         case eval' env v of
@@ -96,7 +96,7 @@ eval !ctx !env0 = convertRuntimeNodes . eval' env0
         k -> nodeFromInteger (mod (integerFromNode (eval' env l)) k)
     applyBuiltin _ env OpIntLt [l, r] = nodeFromBool (integerFromNode (eval' env l) < integerFromNode (eval' env r))
     applyBuiltin _ env OpIntLe [l, r] = nodeFromBool (integerFromNode (eval' env l) <= integerFromNode (eval' env r))
-    applyBuiltin _ env OpEq [l, r] = nodeFromBool (eval' env l == eval' env r)
+    applyBuiltin _ env OpEq [l, r] = nodeFromBool (structEq (eval' env l) (eval' env r))
     applyBuiltin _ env OpTrace [msg, x] = Debug.trace (printNode (eval' env msg)) (eval' env x)
     applyBuiltin _ env OpFail [msg] =
       Exception.throw (EvalError (fromString ("failure: " ++ printNode (eval' env msg))) Nothing)
@@ -152,7 +152,7 @@ hEvalIO hin hout ctx env node =
 evalIO :: IdentContext -> Env -> Node -> IO Node
 evalIO = hEvalIO stdin stdout
 
--- Catch EvalError and convert it to CoreError. Needs a default location in case
+-- | Catch EvalError and convert it to CoreError. Needs a default location in case
 -- no location is available in EvalError.
 catchEvalError :: Location -> a -> IO (Either CoreError a)
 catchEvalError loc a =

src/Juvix/Compiler/Core/Extra.hs

@@ -3,11 +3,13 @@ module Juvix.Compiler.Core.Extra
     module Juvix.Compiler.Core.Extra.Base,
     module Juvix.Compiler.Core.Extra.Recursors,
     module Juvix.Compiler.Core.Extra.Info,
+    module Juvix.Compiler.Core.Extra.Equality,
   )
 where
 
 import Data.HashSet qualified as HashSet
 import Juvix.Compiler.Core.Extra.Base
+import Juvix.Compiler.Core.Extra.Equality
 import Juvix.Compiler.Core.Extra.Info
 import Juvix.Compiler.Core.Extra.Recursors
 import Juvix.Compiler.Core.Language
@@ -92,7 +94,7 @@ convertClosures = umap go
   where
     go :: Node -> Node
     go n = case n of
-      Closure i env b -> substEnv env (mkLambda i b)
+      Closure env (Lambda i b) -> substEnv env (mkLambda i b)
       _ -> n
 
 convertRuntimeNodes :: Node -> Node

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

@@ -93,7 +93,7 @@ mkTypeConstr' = mkTypeConstr Info.empty
 {------------------------------------------------------------------------}
 {- functions on Type -}
 
--- unfold a type into the target and the arguments (left-to-right)
+-- | Unfold a type into the target and the arguments (left-to-right)
 unfoldType' :: Type -> (Type, [(Info, Type)])
 unfoldType' ty = case ty of
   NPi (Pi i l r) -> let (tgt, args) = unfoldType' r in (tgt, (i, l) : args)
@@ -138,7 +138,7 @@ unfoldLambdas = go []
 unfoldLambdas' :: Node -> (Int, Node)
 unfoldLambdas' = first length . unfoldLambdas
 
--- `NodeDetails` is a convenience datatype which provides the most commonly needed
+-- | `NodeDetails` is a convenience datatype which provides the most commonly needed
 -- information about a node in a generic fashion.
 data NodeDetails = NodeDetails
   { -- `nodeInfo` is the info associated with the node,
@@ -212,13 +212,13 @@ destruct = \case
     NodeDetails i [] [] [] (\i' _ -> mkUniv i' l)
   NTyp (TypeConstr i sym args) ->
     NodeDetails i args (map (const 0) args) (map (const Nothing) args) (`mkTypeConstr` sym)
-  Closure i env b ->
+  Closure env (Lambda i b) ->
     NodeDetails
       i
       (b : env)
       (1 : map (const 0) env)
       (fetchBinderInfo i : map (const Nothing) env)
-      (\i' args' -> Closure i' (tl args') (hd args'))
+      (\i' args' -> Closure (tl args') (Lambda i' (hd args')))
   where
     fetchBinderInfo :: Info -> Maybe [BinderInfo]
     fetchBinderInfo i = case Info.lookup kBinderInfo i of

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

@@ -0,0 +1,16 @@
+module Juvix.Compiler.Core.Extra.Equality where
+
+import Juvix.Compiler.Core.Language
+
+-- | Structural equality on Nodes. Semantically the same as `==` but slightly
+-- optimised. Used mostly in the evaluator.
+structEq :: Node -> Node -> Bool
+structEq (NCst (Constant _ v1)) (NCst (Constant _ v2)) = v1 == v2
+structEq (NCtr (Constr _ tag1 args1)) (NCtr (Constr _ tag2 args2)) =
+  tag1 == tag2 && argsEq args1 args2
+  where
+    argsEq :: [Node] -> [Node] -> Bool
+    argsEq [] [] = True
+    argsEq (x:xs) (y:ys) | structEq x y = argsEq xs ys
+    argsEq _ _ = False
+structEq x y = x == y

src/Juvix/Compiler/Core/Language.hs

@@ -18,31 +18,31 @@ import Juvix.Compiler.Core.Language.Base
 {---------------------------------------------------------------------------------}
 {- Program tree datatype -}
 
--- De Bruijn index of a locally bound variable.
+-- | De Bruijn index of a locally bound variable.
 data Var = Var {_varInfo :: !Info, _varIndex :: !Index}
 
--- Global identifier of a function (with corresponding `Node` in the global
+-- | Global identifier of a function (with corresponding `Node` in the global
 -- context).
 data Ident = Ident {_identInfo :: !Info, _identSymbol :: !Symbol}
 
 data Constant = Constant {_constantInfo :: !Info, _constantValue :: !ConstantValue}
 
 data App = App {_appInfo :: !Info, _appLeft :: !Node, _appRight :: !Node}
 
--- A builtin application. A builtin has no corresponding Node. It is treated
+-- | A builtin application. A builtin has no corresponding Node. It is treated
 -- specially by the evaluator and the code generator. For example, basic
 -- arithmetic operations go into `Builtin`. The number of arguments supplied
--- must be equal to the number of arguments expected by the builtin
--- operation (this simplifies evaluation and code generation). If you need
--- partial application, eta-expand with lambdas, e.g., eta-expand `(+) 2` to
--- `\x -> (+) 2 x`. See Transformation/Eta.hs.
+-- must be equal to the number of arguments expected by the builtin operation
+-- (this simplifies evaluation and code generation). If you need partial
+-- application, eta-expand with lambdas, e.g., eta-expand `(+) 2` to `\x -> (+)
+-- 2 x`. See Transformation/Eta.hs.
 data BuiltinApp = BuiltinApp
   { _builtinAppInfo :: !Info,
     _builtinAppOp :: !BuiltinOp,
     _builtinAppArgs :: ![Node]
   }
 
--- A data constructor application. The number of arguments supplied must be
+-- | A data constructor application. The number of arguments supplied must be
 -- equal to the number of arguments expected by the constructor.
 data Constr = Constr
   { _constrInfo :: !Info,
@@ -52,11 +52,11 @@ data Constr = Constr
 
 data Lambda = Lambda {_lambdaInfo :: !Info, _lambdaBody :: !Node}
 
--- `let x := value in body` is not reducible to lambda + application for the purposes
--- of ML-polymorphic / dependent type checking or code generation!
+-- | `let x := value in body` is not reducible to lambda + application for the
+-- purposes of ML-polymorphic / dependent type checking or code generation!
 data Let = Let {_letInfo :: !Info, _letValue :: !Node, _letBody :: !Node}
 
--- One-level case matching on the tag of a data constructor: `Case value
+-- | 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
   { _caseInfo :: !Info,
@@ -65,24 +65,24 @@ data Case = Case
     _caseDefault :: !(Maybe Node)
   }
 
--- Dependent Pi-type. Compilation-time only. Pi implicitly introduces a binder
+-- | Dependent Pi-type. Compilation-time only. Pi implicitly introduces a binder
 -- in the body, exactly like Lambda. So `Pi info ty body` is `Pi x : ty .
 -- body` in more familiar notation, but references to `x` in `body` are via de
 -- Bruijn index. For example, Pi A : Type . A -> A translates to (omitting
 -- Infos): Pi (Univ level) (Pi (Var 0) (Var 1)).
 data Pi = Pi {_piInfo :: !Info, _piType :: !Type, _piBody :: !Type}
 
--- Universe. Compilation-time only.
+-- | Universe. Compilation-time only.
 data Univ = Univ {_univInfo :: !Info, _univLevel :: !Int}
 
--- Type constructor application. Compilation-time only.
+-- | Type constructor application. Compilation-time only.
 data TypeConstr = TypeConstr
   { _typeConstrInfo :: !Info,
     _typeConstrSymbol :: !Symbol,
     _typeConstrArgs :: ![Type]
   }
 
--- `Node` is the type of nodes in the program tree. The nodes themselves
+-- | `Node` is the type of nodes in the program tree. The nodes themselves
 -- contain only runtime-relevant information. Runtime-irrelevant annotations
 -- (including all type information) are stored in the infos associated with each
 -- node.
@@ -101,10 +101,10 @@ data Node
   | NTyp {-# UNPACK #-} !TypeConstr
   | -- Evaluation only: `Closure env body`
     Closure
-      { _closureInfo :: !Info,
-        _closureEnv :: !Env,
-        _closureBody :: !Node
+      { _closureEnv :: !Env,
+        _closureLambda :: {-# UNPACK #-} !Lambda
       }
+  deriving stock (Eq)
 
 -- Other things we might need in the future:
 -- - laziness annotations (converting these to closure/thunk creation should be
@@ -118,7 +118,7 @@ data ConstantValue
 -- Other things we might need in the future:
 -- - ConstFloat or ConstFixedPoint
 
--- `CaseBranch tag argsNum branch`
+-- | `CaseBranch tag argsNum branch`
 -- - `argsNum` is the number of arguments of the constructor tagged with `tag`,
 --   equal to the number of implicit binders above `branch`
 data CaseBranch = CaseBranch {_caseTag :: !Tag, _caseBindersNum :: !Int, _caseBranch :: !Node}
@@ -140,43 +140,100 @@ type Env = [Node]
 
 type Type = Node
 
+instance HasAtomicity Var where
+  atomicity _ = Atom
+
+instance HasAtomicity Ident where
+  atomicity _ = Atom
+
+instance HasAtomicity Constant where
+  atomicity _ = Atom
+
+instance HasAtomicity App where
+  atomicity _ = Aggregate appFixity
+
+instance HasAtomicity BuiltinApp where
+  atomicity BuiltinApp {..}
+    | null _builtinAppArgs = Atom
+    | otherwise = Aggregate lambdaFixity
+
+instance HasAtomicity Constr where
+  atomicity Constr {..}
+    | null _constrArgs = Atom
+    | otherwise = Aggregate lambdaFixity
+
+instance HasAtomicity Lambda where
+  atomicity _ = Aggregate lambdaFixity
+
+instance HasAtomicity Let where
+  atomicity _ = Aggregate lambdaFixity
+
+instance HasAtomicity Case where
+  atomicity _ = Aggregate lambdaFixity
+
+instance HasAtomicity Pi where
+  atomicity _ = Aggregate lambdaFixity
+
+instance HasAtomicity Univ where
+  atomicity _ = Atom
+
+instance HasAtomicity TypeConstr where
+  atomicity _ = Aggregate lambdaFixity
+
 instance HasAtomicity Node where
   atomicity = \case
-    NVar {} -> Atom
-    NIdt {} -> Atom
-    NCst {} -> Atom
-    NApp {} -> Aggregate appFixity
-    NBlt BuiltinApp {..} | null _builtinAppArgs -> Atom
-    NBlt BuiltinApp {} -> Aggregate lambdaFixity
-    NCtr Constr {..} | null _constrArgs -> Atom
-    NCtr Constr {} -> Aggregate lambdaFixity
-    NLam Lambda {} -> Aggregate lambdaFixity
-    NLet Let {} -> Aggregate lambdaFixity
-    NCase Case {} -> Aggregate lambdaFixity
-    NPi Pi {} -> Aggregate lambdaFixity
-    NUniv Univ {} -> Atom
-    NTyp TypeConstr {} -> Aggregate appFixity
+    NVar x -> atomicity x
+    NIdt x -> atomicity x
+    NCst x -> atomicity x
+    NApp x -> atomicity x
+    NBlt x -> atomicity x
+    NCtr x -> atomicity x
+    NLam x -> atomicity x
+    NLet x -> atomicity x
+    NCase x -> atomicity x
+    NPi x -> atomicity x
+    NUniv x -> atomicity x
+    NTyp x -> atomicity x
     Closure {} -> Aggregate lambdaFixity
 
 lambdaFixity :: Fixity
 lambdaFixity = Fixity (PrecNat 0) (Unary AssocPostfix)
 
-instance Eq Node where
-  (==) :: Node -> Node -> Bool
-  NVar (Var _ idx1) == NVar (Var _ idx2) = idx1 == idx2
-  NIdt (Ident _ sym1) == NIdt (Ident _ sym2) = sym1 == sym2
-  NCst (Constant _ v1) == NCst (Constant _ v2) = v1 == v2
-  NApp (App _ l1 r1) == NApp (App _ l2 r2) = l1 == l2 && r1 == r2
-  NBlt (BuiltinApp _ op1 args1) == NBlt (BuiltinApp _ op2 args2) = op1 == op2 && args1 == args2
-  NCtr (Constr _ tag1 args1) == NCtr (Constr _ tag2 args2) = tag1 == tag2 && args1 == args2
-  NLam (Lambda _ b1) == NLam (Lambda _ b2) = b1 == b2
-  NLet (Let _ v1 b1) == NLet (Let _ v2 b2) = v1 == v2 && b1 == b2
-  NCase (Case _ v1 bs1 def1) == NCase (Case _ v2 bs2 def2) = v1 == v2 && bs1 == bs2 && def1 == def2
-  NPi (Pi _ ty1 b1) == NPi (Pi _ ty2 b2) = ty1 == ty2 && b1 == b2
-  NUniv (Univ _ l1) == NUniv (Univ _ l2) = l1 == l2
-  NTyp (TypeConstr _ sym1 args1) == NTyp (TypeConstr _ sym2 args2) = sym1 == sym2 && args1 == args2
-  Closure _ env1 b1 == Closure _ env2 b2 = env1 == env2 && b1 == b2
-  _ == _ = False
+instance Eq Var where
+  (Var _ idx1) == (Var _ idx2) = idx1 == idx2
+
+instance Eq Ident where
+  (Ident _ sym1) == (Ident _ sym2) = sym1 == sym2
+
+instance Eq Constant where
+  (Constant _ v1) == (Constant _ v2) = v1 == v2
+
+instance Eq App where
+  (App _ l1 r1) == (App _ l2 r2) = l1 == l2 && r1 == r2
+
+instance Eq BuiltinApp where
+  (BuiltinApp _ op1 args1) == (BuiltinApp _ op2 args2) = op1 == op2 && args1 == args2
+
+instance Eq Constr where
+  (Constr _ tag1 args1) == (Constr _ tag2 args2) = tag1 == tag2 && args1 == args2
+
+instance Eq Lambda where
+  (Lambda _ b1) == (Lambda _ b2) = b1 == b2
+
+instance Eq Let where
+  (Let _ v1 b1) == (Let _ v2 b2) = v1 == v2 && b1 == b2
+
+instance Eq Case where
+  (Case _ v1 bs1 def1) == (Case _ v2 bs2 def2) = v1 == v2 && bs1 == bs2 && def1 == def2
+
+instance Eq Pi where
+  (Pi _ ty1 b1) == (Pi _ ty2 b2) = ty1 == ty2 && b1 == b2
+
+instance Eq Univ where
+  (Univ _ l1) == (Univ _ l2) = l1 == l2
+
+instance Eq TypeConstr where
+  (TypeConstr _ sym1 args1) == (TypeConstr _ sym2 args2) = sym1 == sym2 && args1 == args2
 
 makeLenses ''Var
 makeLenses ''Ident

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

@@ -155,8 +155,8 @@ instance PrettyCode Node where
           Just ni -> ppCode (ni ^. NameInfo.infoName)
           Nothing -> return $ kwUnnamedIdent <> pretty _typeConstrSymbol
       return $ foldl' (<+>) n' args'
-    Closure {..} ->
-      ppCode (substEnv _closureEnv (mkLambda _closureInfo _closureBody))
+    Closure env l@Lambda {} ->
+      ppCode (substEnv env (NLam l))
 
 instance PrettyCode a => PrettyCode (NonEmpty a) where
   ppCode x = do