Add dependent defaults for the new typechecker (#2541)

This pr adds support for default values that depend on previous
arguments. For instance, see
[test071](https://github.com/anoma/juvix/blob/ca7d0fa06d7fca30c8d9abe87dfb2afc1b9bd8c9/tests/Compilation/positive/test071.juvix).
- After #2540 is implemented, we'll be able to remove the old type
checker (including the aritychecker).

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

@@ -4,6 +4,7 @@ module Juvix.Compiler.Internal.Data.InfoTable
     extendWithReplExpression,
     lookupConstructor,
     lookupConstructorArgTypes,
+    lookupFunctionMaybe,
     lookupFunction,
     lookupConstructorReturnType,
     lookupInductive,
@@ -169,7 +170,7 @@ lookupConstructor f = do
         $ "impossible: "
           <> ppTrace f
           <> " is not in the InfoTable\n"
-          <> "The registered constructors are: "
+          <> "\nThe registered constructors are: "
           <> ppTrace (HashMap.keys tbl)
 
 lookupConstructorArgTypes :: (Member (Reader InfoTable) r) => Name -> Sem r ([InductiveParameter], [Expression])
@@ -188,13 +189,16 @@ lookupInductive f = do
         $ "impossible: "
           <> ppTrace f
           <> " is not in the InfoTable\n"
-          <> "The registered inductives are: "
+          <> "\nThe registered inductives are: "
           <> ppTrace (HashMap.keys tbl)
 
+lookupFunctionMaybe :: forall r. (Member (Reader InfoTable) r) => Name -> Sem r (Maybe FunctionInfo)
+lookupFunctionMaybe f = HashMap.lookup f <$> asks (^. infoFunctions)
+
 lookupFunction :: forall r. (Member (Reader InfoTable) r) => Name -> Sem r FunctionInfo
 lookupFunction f = do
   err <- impossibleErr
-  HashMap.lookupDefault err f <$> asks (^. infoFunctions)
+  fromMaybe err <$> lookupFunctionMaybe f
   where
     impossibleErr :: Sem r a
     impossibleErr = do
@@ -205,7 +209,7 @@ lookupFunction f = do
           <> ppTrace f
           <> " is not in the InfoTable\n"
           <> ppTrace (getLoc f)
-          <> "The registered functions are: "
+          <> "\nThe registered functions are: "
           <> ppTrace (HashMap.keys tbl)
 
 lookupAxiom :: (Member (Reader InfoTable) r) => Name -> Sem r AxiomInfo

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

@@ -12,8 +12,6 @@ type Rename = HashMap VarName VarName
 
 type Subs = HashMap VarName Expression
 
-type VarMap = HashMap VarName VarName
-
 data ApplicationArg = ApplicationArg
   { _appArgIsImplicit :: IsImplicit,
     _appArg :: Expression
@@ -126,16 +124,33 @@ instance HasExpressions SimpleLambda where
     pure (SimpleLambda bi' b')
 
 instance HasExpressions FunctionParameter where
-  leafExpressions f (FunctionParameter m i e) = do
-    e' <- leafExpressions f e
-    pure (FunctionParameter m i e')
+  leafExpressions f FunctionParameter {..} = do
+    ty' <- leafExpressions f _paramType
+    pure
+      FunctionParameter
+        { _paramType = ty',
+          _paramName,
+          _paramImplicit
+        }
 
 instance HasExpressions Function where
   leafExpressions f (Function l r) = do
     l' <- leafExpressions f l
     r' <- leafExpressions f r
     pure (Function l' r')
 
+instance HasExpressions ApplicationArg where
+  leafExpressions f ApplicationArg {..} = do
+    arg' <- leafExpressions f _appArg
+    pure
+      ApplicationArg
+        { _appArg = arg',
+          _appArgIsImplicit
+        }
+
+instance (HasExpressions a) => HasExpressions (Maybe a) where
+  leafExpressions = _Just . leafExpressions
+
 instance HasExpressions Application where
   leafExpressions f (Application l r i) = do
     l' <- leafExpressions f l
@@ -292,16 +307,24 @@ inductiveTypeVarsAssoc def l
     vars :: [VarName]
     vars = def ^.. inductiveParameters . each . inductiveParamName
 
-substitutionApp :: forall r expr. (Member NameIdGen r, HasExpressions expr) => (Maybe Name, Expression) -> expr -> Sem r expr
+substitutionApp :: (Maybe Name, Expression) -> Subs
 substitutionApp (mv, ty) = case mv of
-  Nothing -> return
-  Just v -> substitutionE (HashMap.singleton v ty)
+  Nothing -> mempty
+  Just v -> HashMap.singleton v ty
 
 localsToSubsE :: LocalVars -> Subs
 localsToSubsE l = ExpressionIden . IdenVar <$> l ^. localTyMap
 
+subsKind :: [Name] -> NameKind -> Subs
+subsKind uids k =
+  HashMap.fromList
+    [ (s, toExpression s') | s <- uids, let s' = toExpression (set nameKind k s)
+    ]
+
 substitutionE :: forall r expr. (Member NameIdGen r, HasExpressions expr) => Subs -> expr -> Sem r expr
-substitutionE m = leafExpressions goLeaf
+substitutionE m
+  | null m = pure
+  | otherwise = leafExpressions goLeaf
   where
     goLeaf :: Expression -> Sem r Expression
     goLeaf = \case
@@ -382,6 +405,7 @@ foldApplication f args = case nonEmpty args of
 unfoldApplication' :: Application -> (Expression, NonEmpty ApplicationArg)
 unfoldApplication' (Application l' r' i') = second (|: (ApplicationArg i' r')) (unfoldExpressionApp l')
 
+-- TODO make it tail recursive
 unfoldExpressionApp :: Expression -> (Expression, [ApplicationArg])
 unfoldExpressionApp = \case
   ExpressionApplication (Application l r i) ->

src/Juvix/Compiler/Internal/Language.hs

@@ -391,7 +391,7 @@ instance HasAtomicity SimpleLambda where
   atomicity = const Atom
 
 instance HasAtomicity Let where
-  atomicity l = atomicity (l ^. letExpression)
+  atomicity = const (Aggregate letFixity)
 
 instance HasAtomicity Literal where
   atomicity = \case

src/Juvix/Compiler/Internal/Translation/FromConcrete.hs

@@ -743,7 +743,7 @@ goExpression ::
 goExpression = \case
   ExpressionIdentifier nt -> return (goIden nt)
   ExpressionParensIdentifier nt -> return (goIden nt)
-  ExpressionApplication a -> Internal.ExpressionApplication <$> goApplication a
+  ExpressionApplication a -> goApplication a
   ExpressionCase a -> Internal.ExpressionCase <$> goCase a
   ExpressionNewCase a -> Internal.ExpressionCase <$> goNewCase a
   ExpressionInfixApplication ia -> Internal.ExpressionApplication <$> goInfix ia
@@ -759,35 +759,32 @@ goExpression = \case
   ExpressionHole h -> return (Internal.ExpressionHole h)
   ExpressionInstanceHole h -> return (Internal.ExpressionInstanceHole (fromHole h))
   ExpressionIterator i -> goIterator i
-  ExpressionNamedApplication i -> goNamedApplication i
-  ExpressionNamedApplicationNew i -> goNamedApplicationNew i
+  ExpressionNamedApplication i -> goNamedApplication i []
+  ExpressionNamedApplicationNew i -> goNamedApplicationNew i []
   ExpressionRecordUpdate i -> goRecordUpdateApp i
   ExpressionParensRecordUpdate i -> Internal.ExpressionLambda <$> goRecordUpdate (i ^. parensRecordUpdate)
   where
-    goNamedApplication :: Concrete.NamedApplication 'Scoped -> Sem r Internal.Expression
-    goNamedApplication w = do
+    goNamedApplication :: Concrete.NamedApplication 'Scoped -> [Internal.ApplicationArg] -> Sem r Internal.Expression
+    goNamedApplication w extraArgs = do
       s <- ask @NameSignatures
-      runReader s (runNamedArguments w) >>= goDesugaredNamedApplication
+      runReader s (runNamedArguments w extraArgs) >>= goDesugaredNamedApplication
 
-    goNamedApplicationNew :: Concrete.NamedApplicationNew 'Scoped -> Sem r Internal.Expression
-    goNamedApplicationNew napp = case nonEmpty (napp ^. namedApplicationNewArguments) of
-      Nothing -> return $ goIden (napp ^. namedApplicationNewName)
+    goNamedApplicationNew :: Concrete.NamedApplicationNew 'Scoped -> [Internal.ApplicationArg] -> Sem r Internal.Expression
+    goNamedApplicationNew napp extraArgs = case nonEmpty (napp ^. namedApplicationNewArguments) of
+      Nothing -> return (goIden (napp ^. namedApplicationNewName))
       Just appargs -> do
         let name = napp ^. namedApplicationNewName . scopedIdenName
         sig <- fromJust <$> asks @NameSignatures (^. at (name ^. S.nameId))
         cls <- goArgs appargs
         let args :: [Internal.Name] = appargs ^.. each . namedArgumentNewFunDef . signName . to goSymbol
             -- changes the kind from Variable to Function
-            updateKind :: Internal.Subs =
-              HashMap.fromList
-                [ (s, Internal.toExpression s') | s <- args, let s' = Internal.IdenFunction (set Internal.nameKind KNameFunction s)
-                ]
+            updateKind :: Internal.Subs = Internal.subsKind args KNameFunction
             napp' =
               Concrete.NamedApplication
                 { _namedAppName = napp ^. namedApplicationNewName,
                   _namedAppArgs = nonEmpty' (createArgumentBlocks (sig ^. nameSignatureArgs))
                 }
-        e <- goNamedApplication napp'
+        e <- goNamedApplication napp' extraArgs
         let l =
               Internal.Let
                 { _letClauses = cls,
@@ -841,18 +838,16 @@ goExpression = \case
     goDesugaredNamedApplication :: DesugaredNamedApplication -> Sem r Internal.Expression
     goDesugaredNamedApplication a = do
       let fun = goScopedIden (a ^. dnamedAppIdentifier)
-          updateKind :: Internal.Subs =
-            HashMap.fromList
-              [ (s, Internal.ExpressionIden s') | s <- a ^.. dnamedAppArgs . each . argName . to goSymbol, let s' = Internal.IdenFunction (set Internal.nameKind KNameFunction s)
-              ]
+          updateKind :: Internal.Subs = Internal.subsKind (a ^.. dnamedAppArgs . each . argName . to goSymbol) KNameFunction
           mkAppArg :: Arg -> Internal.ApplicationArg
           mkAppArg arg =
             Internal.ApplicationArg
               { _appArgIsImplicit = arg ^. argImplicit,
                 _appArg = Internal.toExpression (goSymbol (arg ^. argName))
               }
-          argNames :: NonEmpty Internal.ApplicationArg = mkAppArg <$> a ^. dnamedAppArgs
-          app = Internal.foldApplication (Internal.toExpression fun) (toList argNames)
+          namedArgNames :: NonEmpty Internal.ApplicationArg = mkAppArg <$> a ^. dnamedAppArgs
+          allArgs = toList namedArgNames <> a ^. dnamedExtraArgs
+          app = Internal.foldApplication (Internal.toExpression fun) allArgs
       clauses <- mapM mkClause (a ^. dnamedAppArgs)
       expr <-
         Internal.substitutionE updateKind $
@@ -1004,11 +999,26 @@ goExpression = \case
             e' <- goExpression e
             return (Internal.ApplicationArg i e')
 
-    goApplication :: Application -> Sem r Internal.Application
-    goApplication (Application l arg) = do
-      l' <- goExpression l
-      Internal.ApplicationArg {..} <- goApplicationArg arg
-      return (Internal.Application l' _appArg _appArgIsImplicit)
+    goApplication :: Application -> Sem r Internal.Expression
+    goApplication a = do
+      let (f, args) = unfoldApp a
+      args' <- toList <$> mapM goApplicationArg args
+      case f of
+        ExpressionNamedApplication n -> goNamedApplication n args'
+        ExpressionNamedApplicationNew n -> goNamedApplicationNew n args'
+        _ -> do
+          f' <- goExpression f
+          return (Internal.foldApplication f' args')
+      where
+        unfoldApp :: Application -> (Expression, NonEmpty Expression)
+        unfoldApp (Application l' r') =
+          let (f, largs) = go [] l'
+           in (f, largs |: r')
+          where
+            go :: [Expression] -> Expression -> (Expression, [Expression])
+            go ac = \case
+              ExpressionApplication (Application l r) -> go (r : ac) l
+              e -> (e, ac)
 
     goPostfix :: PostfixApplication -> Sem r Internal.Application
     goPostfix (PostfixApplication l op) = do

src/Juvix/Compiler/Internal/Translation/FromConcrete/NamedArguments.hs

@@ -5,6 +5,7 @@ module Juvix.Compiler.Internal.Translation.FromConcrete.NamedArguments
     DesugaredNamedApplication,
     dnamedAppIdentifier,
     dnamedAppArgs,
+    dnamedExtraArgs,
     Arg,
     argName,
     argImplicit,
@@ -18,8 +19,8 @@ import Data.HashMap.Strict qualified as HashMap
 import Data.IntMap.Strict qualified as IntMap
 import Juvix.Compiler.Concrete.Data.ScopedName qualified as S
 import Juvix.Compiler.Concrete.Extra (symbolParsed)
-import Juvix.Compiler.Concrete.Language
 import Juvix.Compiler.Concrete.Translation.FromParsed.Analysis.Scoping.Error
+import Juvix.Compiler.Internal.Extra.Base qualified as Internal
 import Juvix.Prelude
 
 type NameSignatures = HashMap NameId (NameSignature 'Scoped)
@@ -42,7 +43,8 @@ data Arg = Arg
 -- | The result of desugaring a named application
 data DesugaredNamedApplication = DesugaredNamedApplication
   { _dnamedAppIdentifier :: ScopedIden,
-    _dnamedAppArgs :: NonEmpty Arg
+    _dnamedAppArgs :: NonEmpty Arg,
+    _dnamedExtraArgs :: [Internal.ApplicationArg]
   }
 
 makeLenses ''BuilderState
@@ -53,17 +55,22 @@ runNamedArguments ::
   forall r.
   (Members '[NameIdGen, Error ScoperError, Reader NameSignatures] r) =>
   NamedApplication 'Scoped ->
+  [Internal.ApplicationArg] ->
   Sem r DesugaredNamedApplication
-runNamedArguments napp = do
+runNamedArguments napp extraArgs = do
   iniSt <- mkIniBuilderState
-  _dnamedAppArgs <-
+  namedArgs <-
     fmap nonEmpty'
       . execOutputList
       . mapError ErrNamedArgumentsError
       . execState iniSt
       $ helper (getLoc napp)
-  let _dnamedAppIdentifier = napp ^. namedAppName
-  return DesugaredNamedApplication {..}
+  return
+    DesugaredNamedApplication
+      { _dnamedAppIdentifier = napp ^. namedAppName,
+        _dnamedAppArgs = namedArgs,
+        _dnamedExtraArgs = extraArgs
+      }
   where
     mkIniBuilderState :: Sem r BuilderState
     mkIniBuilderState = do

src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/TypeChecking/Checker.hs

@@ -915,7 +915,7 @@ inferExpression' hint e = case e of
                         <> ppTrace (Application l r iapp)
                     )
                 )
-              ty <- substitutionApp (paraName, r') funR
+              ty <- substitutionE (substitutionApp (paraName, r')) funR
               return
                 TypedExpression
                   { _typedExpression =