Add lambda type info (#1845)

This pr adds the `_lambdaType :: Maybe Expression` to `Internal.Lambda`.
This field will be `Nothing` before typechecking and `Just _` after it.

The type is printed if present. For example if the input of `juvix dev
internal typecheck M.juvix --print-result` is the following:
```
id : {A : Type} → A → A
id := λ { a := a}
```
Then the output is as follows:
```
id : {A : Type} → A → A
id {_ω209} := λ : _ω209 → _ω209 {| a := a}
```

src/Juvix/Compiler/Internal/Extra.hs

@@ -101,7 +101,10 @@ mkConcreteType = fmap ConcreteType . go
         b' <- go b
         ty' <- go ty
         return (ExpressionSimpleLambda (SimpleLambda v ty' b'))
-      ExpressionLambda (Lambda c) -> ExpressionLambda . Lambda <$> mapM goClause c
+      ExpressionLambda (Lambda c ty) -> do
+        let _lambdaType = ty >>= go
+        _lambdaClauses <- mapM goClause c
+        return (ExpressionLambda Lambda {..})
       ExpressionFunction (Function l r) -> do
         l' <- goParam l
         r' <- go r
@@ -201,7 +204,10 @@ mkPolyType = fmap PolyType . go
       ExpressionIden IdenConstructor {} -> return t
       ExpressionIden IdenAxiom {} -> return t
       ExpressionIden IdenVar {} -> return t
-      ExpressionLambda (Lambda c) -> ExpressionLambda . Lambda <$> mapM goClause c
+      ExpressionLambda (Lambda c ty) -> do
+        let _lambdaType = ty >>= go
+        _lambdaClauses <- mapM goClause c
+        return (ExpressionLambda Lambda {..})
       ExpressionSimpleLambda (SimpleLambda v ty b) -> do
         b' <- go b
         ty' <- go ty
@@ -219,7 +225,10 @@ instance HasExpressions LambdaClause where
   leafExpressions f (LambdaClause ps b) = LambdaClause ps <$> leafExpressions f b
 
 instance HasExpressions Lambda where
-  leafExpressions f = traverseOf lambdaClauses (traverse (leafExpressions f))
+  leafExpressions f l = do
+    _lambdaClauses <- traverse (leafExpressions f) (l ^. lambdaClauses)
+    _lambdaType <- traverse (leafExpressions f) (l ^. lambdaType)
+    pure Lambda {..}
 
 instance HasExpressions Expression where
   leafExpressions f e = case e of

src/Juvix/Compiler/Internal/Language.hs

@@ -162,8 +162,10 @@ data Case = Case
 
 instance Hashable Case
 
-newtype Lambda = Lambda
-  { _lambdaClauses :: NonEmpty LambdaClause
+data Lambda = Lambda
+  { _lambdaClauses :: NonEmpty LambdaClause,
+    -- | The typechecker fills this field
+    _lambdaType :: Maybe Expression
   }
   deriving stock (Eq, Generic, Data)
 

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

@@ -107,7 +107,8 @@ instance PrettyCode LambdaClause where
 instance PrettyCode Lambda where
   ppCode Lambda {..} = do
     lambdaClauses' <- ppPipeBlock _lambdaClauses
-    return $ kwLambda <+> lambdaClauses'
+    lambdaType' <- mapM ppCode _lambdaType
+    return $ kwLambda <+> (fmap (kwColon <+>) lambdaType') <?+> braces lambdaClauses'
 
 instance PrettyCode a => PrettyCode (WithLoc a) where
   ppCode = ppCode . (^. withLocParam)
@@ -187,7 +188,7 @@ instance PrettyCode FunctionDef where
     clauses' <- mapM ppCode (f ^. funDefClauses)
     return $
       funDefName'
-        <+> kwColonColon
+        <+> kwColon
         <+> funDefType'
           <> line
           <> vsep (toList clauses')
@@ -197,7 +198,7 @@ instance PrettyCode FunctionClause where
     funName <- ppCode (c ^. clauseName)
     clausePatterns' <- hsepMaybe <$> mapM ppCodeAtom (c ^. clausePatterns)
     clauseBody' <- ppCode (c ^. clauseBody)
-    return $ funName <+?> clausePatterns' <+> kwAssign <+> clauseBody'
+    return $ nest 2 (funName <+?> clausePatterns' <+> kwAssign <+> clauseBody')
 
 instance PrettyCode Backend where
   ppCode = \case

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

@@ -288,7 +288,10 @@ goType e = case e of
   Abstract.ExpressionCase {} -> unsupported "case in types"
 
 goLambda :: forall r. (Members '[NameIdGen] r) => Abstract.Lambda -> Sem r Lambda
-goLambda (Abstract.Lambda cl') = Lambda <$> mapM goClause cl'
+goLambda (Abstract.Lambda cl') = do
+  _lambdaClauses <- mapM goClause cl'
+  let _lambdaType :: Maybe Expression = Nothing
+  return Lambda {..}
   where
     goClause :: Abstract.LambdaClause -> Sem r LambdaClause
     goClause (Abstract.LambdaClause ps b) = do

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

@@ -207,10 +207,10 @@ arityCase c = do
 -- types, it is ok to (partially) infer the arity of the lambda from the clause
 -- with the most patterns.
 arityLambda :: Lambda -> Arity
-arityLambda (Lambda cl) =
+arityLambda l =
   foldArity
     UnfoldedArity
-      { _ufoldArityParams = replicate (maximum1 (fmap numPatterns cl)) (ParamExplicit ArityUnknown),
+      { _ufoldArityParams = replicate (maximum1 (fmap numPatterns (l ^. lambdaClauses))) (ParamExplicit ArityUnknown),
         _ufoldArityRest = ArityRestUnknown
       }
   where
@@ -382,7 +382,10 @@ checkLambda ::
   Arity ->
   Lambda ->
   Sem r Lambda
-checkLambda ari (Lambda cl) = Lambda <$> mapM goClause cl
+checkLambda ari l = do
+  let _lambdaType = l ^. lambdaType
+  _lambdaClauses <- mapM goClause (l ^. lambdaClauses)
+  return Lambda {..}
   where
     goClause :: LambdaClause -> Sem r LambdaClause
     goClause (LambdaClause ps b) = do
@@ -409,7 +412,7 @@ checkLambda ari (Lambda cl) = Lambda <$> mapM goClause cl
                   -- TODO. think what to do in this case
                   return (pats, as)
                 (_ : _, []) -> case rest of
-                  ArityRestUnit -> error ("too many patterns in lambda: " <> ppTrace (Lambda cl) <> "\n" <> prettyText ari)
+                  ArityRestUnit -> error ("too many patterns in lambda: " <> ppTrace l <> "\n" <> prettyText ari)
                   ArityRestUnknown -> return (pats, [])
 
 idenArity :: (Members '[Reader LocalVars, Reader InfoTable] r) => Iden -> Sem r Arity

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

@@ -104,7 +104,7 @@ checkStrictlyPositiveOccurrences ty ctorName name recLimit ref =
           helper inside lamBody
 
         helperLambda :: Lambda -> Sem r ()
-        helperLambda (Lambda cl) = mapM_ goClause cl
+        helperLambda l = mapM_ goClause (l ^. lambdaClauses)
           where
             goClause :: LambdaClause -> Sem r ()
             goClause (LambdaClause _ b) = helper inside b

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

@@ -573,11 +573,13 @@ inferExpression' hint e = case e of
           }
 
     goLambda :: Lambda -> Sem r TypedExpression
-    goLambda l@(Lambda cl) = do
+    goLambda l = do
       ty <- case hint of
         Just hi -> return hi
         Nothing -> ExpressionHole <$> freshHole (getLoc l)
-      l' <- Lambda <$> mapM (goClause ty) cl
+      _lambdaClauses <- mapM (goClause ty) (l ^. lambdaClauses)
+      let _lambdaType = Just ty
+          l' = Lambda {..}
       return
         TypedExpression
           { _typedType = ty,

src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/TypeChecking/Data/Inference.hs

@@ -119,7 +119,10 @@ strongNormalize' = go
       return (LambdaClause p b')
 
     goLambda :: Lambda -> Sem r Lambda
-    goLambda (Lambda cl) = Lambda <$> mapM goClause cl
+    goLambda l = do
+      _lambdaClauses <- mapM goClause (l ^. lambdaClauses)
+      _lambdaType <- mapM go (l ^. lambdaType)
+      return Lambda {..}
 
     goSimpleLambda :: SimpleLambda -> Sem r SimpleLambda
     goSimpleLambda (SimpleLambda lamVar lamTy lamBody) = do
@@ -342,9 +345,9 @@ re = reinterpret $ \case
                               _ -> id
                         bicheck (go l1 l2) (local' (go r1 r2))
                     | otherwise = err
-
+                -- NOTE type is ignored, I think it is ok
                 goLambda :: Lambda -> Lambda -> Sem r (Maybe MatchError)
-                goLambda (Lambda l1) (Lambda l2) = case zipExactMay (toList l1) (toList l2) of
+                goLambda (Lambda l1 _) (Lambda l2 _) = case zipExactMay (toList l1) (toList l2) of
                   Just z -> asum <$> mapM (uncurry goClause) z
                   _ -> err
                   where