Improve inference for --new-typechecker

src/Juvix/Compiler/Core/Translation/FromInternal.hs

@@ -15,6 +15,7 @@ import Juvix.Compiler.Core.Translation.FromInternal.Builtins.Nat
 import Juvix.Compiler.Core.Translation.FromInternal.Data
 import Juvix.Compiler.Internal.Data.Name
 import Juvix.Compiler.Internal.Extra qualified as Internal
+import Juvix.Compiler.Internal.Pretty (ppTrace)
 import Juvix.Compiler.Internal.Pretty qualified as Internal
 import Juvix.Compiler.Internal.Translation.Extra qualified as Internal
 import Juvix.Compiler.Internal.Translation.FromInternal.Analysis.TypeChecking qualified as InternalTyped
@@ -856,7 +857,8 @@ goIden i = do
           )
   case i of
     Internal.IdenVar n -> do
-      k <- HashMap.lookupDefault impossible id_ <$> asks (^. indexTableVars)
+      let err = error ("impossible: var not found: " <> ppTrace n <> " at " <> prettyText (getLoc n))
+      k <- HashMap.lookupDefault err id_ <$> asks (^. indexTableVars)
       varsNum <- asks (^. indexTableVarsNum)
       return (mkVar (setInfoLocation (n ^. nameLoc) (Info.singleton (NameInfo (n ^. nameText)))) (varsNum - k - 1))
     Internal.IdenFunction n -> do

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

@@ -43,3 +43,6 @@ addTypeMapping v v' = over localTyMap (HashMap.insert v v')
 
 withEmptyLocalVars :: Sem (Reader LocalVars ': r) a -> Sem r a
 withEmptyLocalVars = runReader emptyLocalVars
+
+withLocalVars :: LocalVars -> Sem (Reader LocalVars ': r) a -> Sem r a
+withLocalVars = runReader

src/Juvix/Compiler/Internal/Extra.hs

@@ -109,7 +109,7 @@ genFieldProjection _funDefName _funDefBuiltin info fieldIx = do
       saturatedTy = unnamedParameter' implicity (constructorReturnType info)
       inductiveArgs = map inductiveToFunctionParam inductiveParams
       retTy = constrArgs !! fieldIx
-  return
+  cloneFunctionDefSameName
     FunctionDef
       { _funDefExamples = [],
         _funDefTerminating = False,
@@ -180,3 +180,47 @@ mkLetClauses pre = goSCC <$> (toList (buildLetMutualBlocks pre))
               getFun :: PreLetStatement -> FunctionDef
               getFun = \case
                 PreLetFunctionDef f -> f
+
+inlineLet :: forall r. (Members '[NameIdGen] r) => Let -> Sem r Expression
+inlineLet l = do
+  (lclauses, subs) <-
+    runOutputList
+      . execState (mempty @Subs)
+      $ forM (l ^. letClauses) helper
+  body' <- substitutionE subs (l ^. letExpression)
+  return $ case nonEmpty lclauses of
+    Nothing -> body'
+    Just cl' ->
+      ExpressionLet
+        Let
+          { _letClauses = cl',
+            _letExpression = body'
+          }
+  where
+    helper :: forall r'. (r' ~ (State Subs ': Output LetClause ': r)) => LetClause -> Sem r' ()
+    helper c = do
+      subs <- get
+      c' <- substitutionE subs c
+      case subsClause c' of
+        Nothing -> output c'
+        Just (n, b) -> modify' @Subs (set (at n) (Just b))
+
+    subsClause :: LetClause -> Maybe (Name, Expression)
+    subsClause = \case
+      LetMutualBlock {} -> Nothing
+      LetFunDef f -> mkAssoc f
+      where
+        mkAssoc :: FunctionDef -> Maybe (Name, Expression)
+        mkAssoc = \case
+          FunctionDef
+            { _funDefType = ExpressionHole {},
+              _funDefBody = body,
+              _funDefName = name,
+              _funDefArgsInfo = []
+            } -> Just (name, body)
+          _ -> Nothing
+
+cloneFunctionDefSameName :: (Members '[NameIdGen] r) => FunctionDef -> Sem r FunctionDef
+cloneFunctionDefSameName f = do
+  f' <- clone f
+  return (set funDefName (f ^. funDefName) f')

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

@@ -305,13 +305,13 @@ substitutionE m = leafExpressions goLeaf
   where
     goLeaf :: Expression -> Sem r Expression
     goLeaf = \case
-      ExpressionIden i -> goIden i
+      ExpressionIden i -> goName (i ^. idenName)
       e -> return e
-    goIden :: Iden -> Sem r Expression
-    goIden i = case i of
-      IdenVar v
-        | Just e <- HashMap.lookup v m -> clone e
-      _ -> return $ ExpressionIden i
+    goName :: Name -> Sem r Expression
+    goName n =
+      case HashMap.lookup n m of
+        Just e -> clone e
+        Nothing -> return (toExpression n)
 
 smallUniverseE :: Interval -> Expression
 smallUniverseE = ExpressionUniverse . SmallUniverse

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

@@ -234,14 +234,15 @@ instance Clonable ArgInfo where
           _argInfoName
         }
 
+-- | Note that the name of the function is fresh. This is desirable when the
+-- functionDef is part of a let.
 instance Clonable FunctionDef where
   freshNameIds :: (Members '[Reader FreshBindersContext, NameIdGen] r) => FunctionDef -> Sem r FunctionDef
   freshNameIds fun@FunctionDef {..} = do
     ty' <- freshNameIds _funDefType
     underBinder fun $ \fun' -> do
       body' <- freshNameIds _funDefBody
       defaultSig' <- freshNameIds _funDefArgsInfo
-
       return
         FunctionDef
           { _funDefName = fun' ^. funDefName,

src/Juvix/Compiler/Internal/Language.hs

@@ -427,7 +427,7 @@ instance HasAtomicity ConstructorApp where
 
 instance HasAtomicity PatternArg where
   atomicity p
-    | Implicit <- p ^. patternArgIsImplicit = Atom
+    | isImplicitOrInstance (p ^. patternArgIsImplicit) = Atom
     | isJust (p ^. patternArgName) = Atom
     | otherwise = atomicity (p ^. patternArgPattern)
 

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

@@ -11,7 +11,8 @@ import Juvix.Compiler.Internal.Data.InstanceInfo (instanceInfoResult, instanceTa
 import Juvix.Compiler.Internal.Data.LocalVars
 import Juvix.Compiler.Internal.Data.NameDependencyInfo
 import Juvix.Compiler.Internal.Data.TypedHole
-import Juvix.Compiler.Internal.Extra
+import Juvix.Compiler.Internal.Extra.Base
+import Juvix.Compiler.Internal.Language
 import Juvix.Compiler.Internal.Pretty.Options
 import Juvix.Compiler.Internal.Translation.FromInternal.Analysis.ArityChecking.Data.Types (Arity)
 import Juvix.Compiler.Internal.Translation.FromInternal.Analysis.TypeChecking.CheckerNew.Arity qualified as New

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

@@ -18,6 +18,7 @@ import Juvix.Compiler.Builtins
 import Juvix.Compiler.Concrete.Data.Scope.Base (ScoperState, scoperScopedConstructorFields, scoperScopedSignatures)
 import Juvix.Compiler.Concrete.Data.ScopedName qualified as S
 import Juvix.Compiler.Concrete.Extra qualified as Concrete
+import Juvix.Compiler.Concrete.Gen qualified as Gen
 import Juvix.Compiler.Concrete.Language qualified as Concrete
 import Juvix.Compiler.Concrete.Translation.FromParsed.Analysis.Scoping qualified as Scoper
 import Juvix.Compiler.Concrete.Translation.FromParsed.Analysis.Scoping.Error
@@ -784,16 +785,17 @@ goExpression = \case
             napp' =
               Concrete.NamedApplication
                 { _namedAppName = napp ^. namedApplicationNewName,
-                  _namedAppArgs = nonEmpty' $ createArgumentBlocks (sig ^. nameSignatureArgs)
+                  _namedAppArgs = nonEmpty' (createArgumentBlocks (sig ^. nameSignatureArgs))
                 }
         e <- goNamedApplication napp'
+        let l =
+              Internal.Let
+                { _letClauses = cls,
+                  _letExpression = e
+                }
         expr <-
-          Internal.substitutionE updateKind
-            . Internal.ExpressionLet
-            $ Internal.Let
-              { _letClauses = cls,
-                _letExpression = e
-              }
+          Internal.substitutionE updateKind l
+            >>= Internal.inlineLet
         Internal.clone expr
         where
           goArgs :: NonEmpty (NamedArgumentNew 'Scoped) -> Sem r (NonEmpty Internal.LetClause)
@@ -803,7 +805,7 @@ goExpression = \case
               goArg = fmap Internal.PreLetFunctionDef . goFunctionDef . (^. namedArgumentNewFunDef)
 
           createArgumentBlocks :: [NameBlock 'Scoped] -> [ArgumentBlock 'Scoped]
-          createArgumentBlocks sblocks = snd $ foldr goBlock (args0, []) sblocks
+          createArgumentBlocks = snd . foldr goBlock (args0, [])
             where
               args0 :: HashSet S.Symbol = HashSet.fromList $ fmap (^. namedArgumentNewFunDef . signName) (toList appargs)
               goBlock :: NameBlock 'Scoped -> (HashSet S.Symbol, [ArgumentBlock 'Scoped]) -> (HashSet S.Symbol, [ArgumentBlock 'Scoped])
@@ -813,11 +815,11 @@ goExpression = \case
                 where
                   namesInBlock =
                     HashSet.intersection
-                      (HashSet.fromList $ HashMap.keys _nameBlock)
+                      (HashSet.fromList (HashMap.keys _nameBlock))
                       (HashSet.map (^. S.nameConcrete) args)
-                  argNames = HashMap.fromList $ map (\n -> (n ^. S.nameConcrete, n)) $ toList args
+                  argNames = HashMap.fromList . map (\n -> (n ^. S.nameConcrete, n)) $ toList args
                   args' = HashSet.filter (not . flip HashSet.member namesInBlock . (^. S.nameConcrete)) args
-                  _argBlockArgs = nonEmpty' $ map goArg (toList namesInBlock)
+                  _argBlockArgs = nonEmpty' (map goArg (toList namesInBlock))
                   block' =
                     ArgumentBlock
                       { _argBlockDelims = Irrelevant Nothing,
@@ -829,11 +831,11 @@ goExpression = \case
                     NamedArgument
                       { _namedArgName = sym,
                         _namedArgAssignKw = Irrelevant dummyKw,
-                        _namedArgValue = Concrete.ExpressionIdentifier $ ScopedIden name Nothing
+                        _namedArgValue = Concrete.ExpressionIdentifier (ScopedIden name Nothing)
                       }
                     where
                       name = over S.nameConcrete NameUnqualified $ fromJust $ HashMap.lookup sym argNames
-                      dummyKw = KeywordRef (asciiKw ":=") dummyLoc Ascii
+                      dummyKw = run (runReader dummyLoc (Gen.kw Gen.kwAssign))
                       dummyLoc = getLoc sym
 
     goDesugaredNamedApplication :: DesugaredNamedApplication -> Sem r Internal.Expression

src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/ArityChecking/Data/Types.hs

@@ -1,6 +1,7 @@
 module Juvix.Compiler.Internal.Translation.FromInternal.Analysis.ArityChecking.Data.Types where
 
-import Juvix.Compiler.Internal.Extra
+import Juvix.Compiler.Internal.Extra.Base
+import Juvix.Compiler.Internal.Language
 import Juvix.Prelude
 import Juvix.Prelude.Pretty
 

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

@@ -559,7 +559,7 @@ checkPattern = go
     go :: FunctionParameter -> PatternArg -> Sem r PatternArg
     go argTy patArg = do
       matchIsImplicit (argTy ^. paramImplicit) patArg
-      tyVarMap <- fmap (ExpressionIden . IdenVar) . (^. localTyMap) <$> get
+      tyVarMap <- localsToSubsE <$> get
       ty <- substitutionE tyVarMap (argTy ^. paramType)
       let pat = patArg ^. patternArgPattern
           name = patArg ^. patternArgName

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

@@ -544,12 +544,25 @@ checkClause clauseLoc clauseType clausePats body = do
     helper :: [PatternArg] -> Expression -> Sem r (LocalVars, ([PatternArg], Expression))
     helper pats ty = runState emptyLocalVars (go pats ty)
 
+    genPatternWildcard :: Interval -> FunctionParameter -> Sem (State LocalVars ': r) PatternArg
+    genPatternWildcard loc par = do
+      let impl = par ^. paramImplicit
+      var <- maybe (varFromWildcard (Wildcard loc)) return (par ^. paramName)
+      addPatternVar var (par ^. paramType) Nothing
+      return
+        PatternArg
+          { _patternArgIsImplicit = impl,
+            _patternArgName = Nothing,
+            _patternArgPattern = PatternVariable var
+          }
+
     go :: [PatternArg] -> Expression -> Sem (State LocalVars ': r) ([PatternArg], Expression)
     go pats bodyTy = case pats of
       [] -> do
         (bodyParams, bodyRest) <- unfoldFunType' bodyTy
-        guessedBodyParams <- unfoldArity <$> guessArity body
-        let pref' :: [IsImplicit] = map (^. paramImplicit) (take pref bodyParams)
+        locals <- get
+        guessedBodyParams <- withLocalVars locals (unfoldArity <$> guessArity body)
+        let pref' :: [FunctionParameter] = take pref bodyParams
             pref :: Int = aI - targetI
             preImplicits = length . takeWhile isImplicitOrInstance
             aI :: Int = preImplicits (map (^. paramImplicit) bodyParams)
@@ -559,10 +572,9 @@ checkClause clauseLoc clauseType clausePats body = do
                 let n = length pref'
                     bodyParams' = drop n bodyParams
                     ty' = foldFunType bodyParams' bodyRest
-                wildcards <- mapM (genWildcard clauseLoc) pref'
+                wildcards <- mapM (genPatternWildcard clauseLoc) pref'
                 return (wildcards, ty')
-            | otherwise -> do
-                return ([], bodyTy)
+            | otherwise -> return ([], bodyTy)
       p : ps -> do
         bodyTy' <- weakNormalize bodyTy
         case bodyTy' of
@@ -581,9 +593,9 @@ checkClause clauseLoc clauseType clausePats body = do
                     loc :: Interval
                     loc = getLoc par
 
-                    insertWildcard :: IsImplicit -> Sem (State LocalVars ': r) ([PatternArg], Expression)
-                    insertWildcard impl = do
-                      w <- genWildcard loc impl
+                    insertWildcard :: Sem (State LocalVars ': r) ([PatternArg], Expression)
+                    insertWildcard = do
+                      w <- genPatternWildcard loc par
                       go (w : p : ps) bodyTy'
 
                 case (p ^. patternArgIsImplicit, par ^. paramImplicit) of
@@ -592,10 +604,10 @@ checkClause clauseLoc clauseType clausePats body = do
                   (ImplicitInstance, ImplicitInstance) -> checkPatternAndContinue
                   (Implicit, Explicit) -> throwWrongIsImplicit p Implicit
                   (ImplicitInstance, Explicit) -> throwWrongIsImplicit p ImplicitInstance
-                  (Explicit, Implicit) -> insertWildcard Implicit
-                  (ImplicitInstance, Implicit) -> insertWildcard Implicit
-                  (Explicit, ImplicitInstance) -> insertWildcard ImplicitInstance
-                  (Implicit, ImplicitInstance) -> insertWildcard ImplicitInstance
+                  (Explicit, Implicit) -> insertWildcard
+                  (ImplicitInstance, Implicit) -> insertWildcard
+                  (Explicit, ImplicitInstance) -> insertWildcard
+                  (Implicit, ImplicitInstance) -> insertWildcard
         where
           throwWrongIsImplicit :: (Members '[Error TypeCheckerError] r') => PatternArg -> IsImplicit -> Sem r' a
           throwWrongIsImplicit patArg expected =
@@ -647,6 +659,12 @@ matchIsImplicit expected actual =
           _wrongPatternIsImplicitActual = actual
         }
 
+addPatternVar :: (Members '[State LocalVars, Inference] r) => VarName -> Expression -> Maybe Name -> Sem r ()
+addPatternVar v ty argName = do
+  modify (addType v ty)
+  registerIdenType v ty
+  whenJust argName (\v' -> modify (addTypeMapping v' v))
+
 checkPattern ::
   forall r.
   (Members '[Reader InfoTable, Error TypeCheckerError, State LocalVars, Inference, NameIdGen, State FunctionsTable] r) =>
@@ -662,9 +680,9 @@ checkPattern = go
       ty <- substitutionE tyVarMap (argTy ^. paramType)
       let pat = patArg ^. patternArgPattern
           name = patArg ^. patternArgName
-      whenJust name (\n -> addVar n ty argTy)
+      whenJust name (\n -> addPatternVar n ty (argTy ^. paramName))
       pat' <- case pat of
-        PatternVariable v -> addVar v ty argTy $> pat
+        PatternVariable v -> addPatternVar v ty (argTy ^. paramName) $> pat
         PatternWildcardConstructor {} -> impossible
         PatternConstructorApp a -> goPatternConstructor pat ty a
       return (set patternArgPattern pat' patArg)
@@ -714,12 +732,6 @@ checkPattern = go
                 )
               PatternConstructorApp <$> goConstr (IdenInductive ind) a tyArgs
 
-        addVar :: VarName -> Expression -> FunctionParameter -> Sem r ()
-        addVar v ty argType = do
-          modify (addType v ty)
-          registerIdenType v ty
-          whenJust (argType ^. paramName) (\v' -> modify (addTypeMapping v' v))
-
         goConstr :: Iden -> ConstructorApp -> [(InductiveParameter, Expression)] -> Sem r ConstructorApp
         goConstr inductivename app@(ConstructorApp c ps _) ctx = do
           (_, psTys) <- constructorArgTypes <$> lookupConstructor c
@@ -1129,7 +1141,7 @@ holesHelper mhint expr = do
       where
         goImplArgs :: Int -> [ApplicationArg] -> Sem r [ApplicationArg]
         goImplArgs 0 as = return as
-        goImplArgs k ((ApplicationArg Implicit _) : as) = goImplArgs (k - 1) as
+        goImplArgs k (ApplicationArg Implicit _ : as) = goImplArgs (k - 1) as
         goImplArgs _ as = return as
 
     goArgs :: forall r'. (r' ~ State AppBuilder ': r) => Sem r' ()
@@ -1432,7 +1444,7 @@ typeArity = weakNormalize >=> go
 
 guessArity ::
   forall r.
-  (Members '[Reader InfoTable, Inference] r) =>
+  (Members '[Reader InfoTable, Inference, Reader LocalVars] r) =>
   Expression ->
   Sem r Arity
 guessArity = \case
@@ -1449,7 +1461,7 @@ guessArity = \case
   ExpressionCase l -> arityCase l
   where
     idenHelper :: Iden -> Sem r Arity
-    idenHelper = withEmptyLocalVars . idenArity
+    idenHelper = idenArity
 
     appHelper :: Application -> Sem r Arity
     appHelper a = do
@@ -1484,7 +1496,7 @@ arityUniverse = ArityUnit
 simplelambda :: a
 simplelambda = error "simple lambda expressions are not supported by the arity checker"
 
-arityLambda :: forall r. (Members '[Reader InfoTable, Inference] r) => Lambda -> Sem r Arity
+arityLambda :: forall r. (Members '[Reader InfoTable, Inference, Reader LocalVars] r) => Lambda -> Sem r Arity
 arityLambda l = do
   aris <- mapM guessClauseArity (l ^. lambdaClauses)
   return $
@@ -1521,13 +1533,13 @@ guessPatternArgArity p =
           }
     }
 
-arityLet :: (Members '[Reader InfoTable, Inference] r) => Let -> Sem r Arity
+arityLet :: (Members '[Reader InfoTable, Inference, Reader LocalVars] r) => Let -> Sem r Arity
 arityLet l = guessArity (l ^. letExpression)
 
 -- | All branches should have the same arity. If they are all the same, we
 -- return that, otherwise we return ArityBlocking. Probably something better can
 -- be done.
-arityCase :: (Members '[Reader InfoTable, Inference] r) => Case -> Sem r Arity
+arityCase :: (Members '[Reader InfoTable, Inference, Reader LocalVars] r) => Case -> Sem r Arity
 arityCase c = do
   aris <- mapM (guessArity . (^. caseBranchExpression)) (c ^. caseBranches)
   return

src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/TypeChecking/CheckerNew/Arity.hs

@@ -1,6 +1,7 @@
 module Juvix.Compiler.Internal.Translation.FromInternal.Analysis.TypeChecking.CheckerNew.Arity where
 
-import Juvix.Compiler.Internal.Extra
+import Juvix.Compiler.Internal.Extra.Base
+import Juvix.Compiler.Internal.Language
 import Juvix.Prelude
 import Juvix.Prelude.Pretty