add NormalizedExpression

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

@@ -362,7 +362,7 @@ goType ::
   Sem r Type
 goType ty = do
   normTy <- InternalTyped.strongNormalize'' ty
-  squashApps <$> goExpression normTy
+  squashApps <$> goExpression (normTy ^. Internal.normalizedExpression)
 
 mkFunBody ::
   forall r.

src/Juvix/Compiler/Internal/Language.hs

@@ -428,6 +428,12 @@ newtype ModuleIndex = ModuleIndex
   }
   deriving stock (Data)
 
+-- | An expression that maybe has been normalized
+data NormalizedExpression = NormalizedExpression
+  { _normalizedExpression :: Expression,
+    _normalizedExpressionOriginal :: Expression
+  }
+
 makeLenses ''ModuleIndex
 makeLenses ''ArgInfo
 makeLenses ''WildcardConstructor
@@ -454,6 +460,10 @@ makeLenses ''FunctionParameter
 makeLenses ''InductiveParameter
 makeLenses ''ConstructorDef
 makeLenses ''ConstructorApp
+makeLenses ''NormalizedExpression
+
+instance HasLoc NormalizedExpression where
+  getLoc = getLoc . (^. normalizedExpressionOriginal)
 
 instance Eq ModuleIndex where
   (==) = (==) `on` (^. moduleIxModule . moduleName)

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

@@ -38,7 +38,7 @@ typeCheckExpressionType exp = do
     . mapError (JuvixError @TypeCheckerError)
     . runInferenceDef
     $ inferExpression Nothing exp
-      >>= traverseOf typedType strongNormalize
+      >>= traverseOf typedType strongNormalize_
 
 typeCheckExpression ::
   (Members '[Error JuvixError, State Artifacts, Termination] r) =>

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

@@ -64,7 +64,7 @@ checkStrictlyPositiveOccurrences ::
   CheckPositivityArgs ->
   Sem r ()
 checkStrictlyPositiveOccurrences p = do
-  typeOfConstr <- strongNormalize (p ^. checkPositivityArgsTypeOfConstructor)
+  typeOfConstr <- strongNormalize_ (p ^. checkPositivityArgsTypeOfConstructor)
   go False typeOfConstr
   where
     indInfo = p ^. checkPositivityArgsInductive

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

@@ -369,12 +369,11 @@ checkInstanceType ::
 checkInstanceType FunctionDef {..} = do
   ty <- strongNormalize _funDefType
   let mi =
-        instanceFromTypedExpression
-          ( TypedExpression
-              { _typedType = ty,
-                _typedExpression = ExpressionIden (IdenFunction _funDefName)
-              }
-          )
+        instanceFromTypedExpression $
+          TypedExpression
+            { _typedType = ty ^. normalizedExpression,
+              _typedExpression = ExpressionIden (IdenFunction _funDefName)
+            }
   case mi of
     Just ii@InstanceInfo {..} -> do
       tab <- ask
@@ -401,10 +400,16 @@ checkInstanceType FunctionDef {..} = do
         _ ->
           throw (ErrNotATrait (NotATrait _paramType))
 
-checkInstanceParam :: (Member (Error TypeCheckerError) r) => InfoTable -> Expression -> Sem r ()
-checkInstanceParam tab ty = case traitFromExpression mempty ty of
+checkInstanceParam ::
+  (Member (Error TypeCheckerError) r) =>
+  InfoTable ->
+  NormalizedExpression ->
+  Sem r ()
+checkInstanceParam tab ty' = case traitFromExpression mempty ty of
   Just InstanceApp {..} | isTrait tab _instanceAppHead -> return ()
   _ -> throw (ErrNotATrait (NotATrait ty))
+  where
+    ty = ty' ^. normalizedExpression
 
 checkCoercionType ::
   forall r.
@@ -416,7 +421,7 @@ checkCoercionType FunctionDef {..} = do
   let mi =
         coercionFromTypedExpression
           ( TypedExpression
-              { _typedType = ty,
+              { _typedType = ty ^. normalizedExpression,
                 _typedExpression = ExpressionIden (IdenFunction _funDefName)
               }
           )
@@ -456,11 +461,16 @@ checkExpression expectedTy e = do
       e' <- strongNormalize e
       inferred' <- strongNormalize (inferred ^. typedType)
       expected' <- strongNormalize expectedTy
+      let thing =
+            WrongTypeThingExpression
+              MkWrongTypeThingExpression
+                { _wrongTypeNormalizedExpression = e',
+                  _wrongTypeInferredExpression = inferred ^. typedExpression
+                }
       throw
         . ErrWrongType
         $ WrongType
-          { _wrongTypeThing = Left e',
-            _wrongTypeThingWithHoles = Just (Left (inferred ^. typedExpression)),
+          { _wrongTypeThing = thing,
             _wrongTypeActual = inferred',
             _wrongTypeExpected = expected'
           }
@@ -497,7 +507,7 @@ checkFunctionParameter FunctionParameter {..} = do
     checkInstanceParam tab ty'
   return
     FunctionParameter
-      { _paramType = ty',
+      { _paramType = ty' ^. normalizedExpression,
         _paramName,
         _paramImplicit
       }
@@ -728,15 +738,13 @@ checkPattern = go
               constrName = a ^. constrAppConstructor
               err :: MatchError -> Sem r ()
               err m =
-                throw
-                  ( ErrWrongType
-                      WrongType
-                        { _wrongTypeThing = Right pat,
-                          _wrongTypeThingWithHoles = Nothing,
-                          _wrongTypeExpected = m ^. matchErrorRight,
-                          _wrongTypeActual = m ^. matchErrorLeft
-                        }
-                  )
+                throw $
+                  ErrWrongType
+                    WrongType
+                      { _wrongTypeThing = WrongTypeThingPattern pat,
+                        _wrongTypeExpected = m ^. matchErrorRight,
+                        _wrongTypeActual = m ^. matchErrorLeft
+                      }
           case s of
             Left hole -> do
               let indParams = info ^. constructorInfoInductiveParameters
@@ -754,15 +762,13 @@ checkPattern = go
             Right (ind, tyArgs) -> do
               when
                 (ind /= constrIndName)
-                ( throw
-                    ( ErrWrongConstructorType
-                        WrongConstructorType
-                          { _wrongCtorTypeName = constrName,
-                            _wrongCtorTypeExpected = ind,
-                            _wrongCtorTypeActual = constrIndName
-                          }
-                    )
-                )
+                $ throw
+                $ ErrWrongConstructorType
+                  WrongConstructorType
+                    { _wrongCtorTypeName = constrName,
+                      _wrongCtorTypeExpected = ind,
+                      _wrongCtorTypeActual = constrIndName
+                    }
               PatternConstructorApp <$> goConstr (IdenInductive ind) a tyArgs
 
         goConstr :: Iden -> ConstructorApp -> [(InductiveParameter, Expression)] -> Sem r ConstructorApp

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

@@ -8,6 +8,7 @@ module Juvix.Compiler.Internal.Translation.FromInternal.Analysis.TypeChecking.Da
     queryMetavar,
     registerIdenType,
     strongNormalize'',
+    strongNormalize_,
     iniState,
     strongNormalize,
     weakNormalize,
@@ -35,8 +36,8 @@ data MetavarState
     Refined Expression
 
 data MatchError = MatchError
-  { _matchErrorLeft :: Expression,
-    _matchErrorRight :: Expression
+  { _matchErrorLeft :: NormalizedExpression,
+    _matchErrorRight :: NormalizedExpression
   }
 
 makeLenses ''MatchError
@@ -46,7 +47,7 @@ data Inference :: Effect where
   QueryMetavar :: Hole -> Inference m (Maybe Expression)
   RegisterIdenType :: Name -> Expression -> Inference m ()
   RememberFunctionDef :: FunctionDef -> Inference m ()
-  StrongNormalize :: Expression -> Inference m Expression
+  StrongNormalize :: Expression -> Inference m NormalizedExpression
   WeakNormalize :: Expression -> Inference m Expression
 
 makeSem ''Inference
@@ -131,9 +132,18 @@ queryMetavarFinal h = do
     Just (ExpressionHole h') -> queryMetavarFinal h'
     _ -> return m
 
+strongNormalize_ :: (Members '[Inference] r) => Expression -> Sem r Expression
+strongNormalize_ = fmap (^. normalizedExpression) . strongNormalize
+
 -- FIXME the returned expression should have the same location as the original
-strongNormalize' :: forall r. (Members '[ResultBuilder, State InferenceState, NameIdGen] r) => Expression -> Sem r Expression
-strongNormalize' = go
+strongNormalize' :: forall r. (Members '[ResultBuilder, State InferenceState, NameIdGen] r) => Expression -> Sem r NormalizedExpression
+strongNormalize' original = do
+  normalized <- go original
+  return
+    NormalizedExpression
+      { _normalizedExpression = normalized,
+        _normalizedExpressionOriginal = original
+      }
   where
     go :: Expression -> Sem r Expression
     go e = case e of
@@ -362,14 +372,30 @@ runInferenceState inis = reinterpret (runState inis) $ \case
               where
                 ok :: Sem r (Maybe MatchError)
                 ok = return Nothing
+
                 check :: Bool -> Sem r (Maybe MatchError)
                 check b
                   | b = ok
                   | otherwise = err
+
                 bicheck :: Sem r (Maybe MatchError) -> Sem r (Maybe MatchError) -> Sem r (Maybe MatchError)
                 bicheck = liftA2 (<|>)
+
+                normalizedB =
+                  NormalizedExpression
+                    { _normalizedExpression = normB,
+                      _normalizedExpressionOriginal = inputB
+                    }
+
+                normalizedA =
+                  NormalizedExpression
+                    { _normalizedExpression = normA,
+                      _normalizedExpressionOriginal = inputA
+                    }
+
                 err :: Sem r (Maybe MatchError)
-                err = return (Just (MatchError normA normB))
+                err = return (Just (MatchError normalizedA normalizedB))
+
                 goHole :: Hole -> Expression -> Sem r (Maybe MatchError)
                 goHole h t = do
                   r <- queryMetavar' h
@@ -382,7 +408,7 @@ runInferenceState inis = reinterpret (runState inis) $ \case
                       | ExpressionHole h' <- holTy, h' == hol = return ()
                       | otherwise =
                           do
-                            holTy' <- strongNormalize' holTy
+                            holTy' <- (^. normalizedExpression) <$> strongNormalize' holTy
                             let er =
                                   ErrUnsolvedMeta
                                     UnsolvedMeta
@@ -392,7 +418,7 @@ runInferenceState inis = reinterpret (runState inis) $ \case
                             when (LeafExpressionHole hol `elem` holTy' ^.. leafExpressions) (throw er)
                             s <- gets (fromJust . (^. inferenceMap . at hol))
                             case s of
-                              Fresh -> modify (over inferenceMap (HashMap.insert hol (Refined holTy')))
+                              Fresh -> modify (set (inferenceMap . at hol) (Just (Refined holTy')))
                               Refined {} -> impossible
 
                 goIden :: Iden -> Iden -> Sem r (Maybe MatchError)
@@ -521,7 +547,7 @@ functionDefEval f = do
   return r
   where
     strongNorm :: (Members '[ResultBuilder, NameIdGen] r) => Expression -> Sem r Expression
-    strongNorm = evalState iniState . strongNormalize'
+    strongNorm = evalState iniState . fmap (^. normalizedExpression) . strongNormalize'
 
     isUniverse :: Expression -> Bool
     isUniverse = \case
@@ -583,7 +609,10 @@ registerFunctionDef :: (Members '[ResultBuilder, Error TypeCheckerError, NameIdG
 registerFunctionDef f = whenJustM (functionDefEval f) $ \e ->
   addFunctionDef (f ^. funDefName) e
 
-strongNormalize'' :: (Members '[Reader FunctionsTable, NameIdGen] r) => Expression -> Sem r Expression
+strongNormalize'' ::
+  (Members '[Reader FunctionsTable, NameIdGen] r) =>
+  Expression ->
+  Sem r NormalizedExpression
 strongNormalize'' ty = do
   ftab <- ask
   let importCtx =