Fix locations in Internal hole substitution (only for the case of substituting identifiers) (#2995)

Type checking messes up the locations by substituting the holes
(instance holes and ordinary holes) without adjusting the location of
the expression substituted into the hole. Instead, the location of the
expression substituted into the hole is preserved. This messes up
locations in type-checked Internal, because the substituted expressions
can come from anywhere. Later on, the error locations are wrong in Core,
and get wrongly displayed e.g. for pattern matching coverage errors.

This PR implements a partial solution for the (most common) case when
the substituted expression is an identifier. In the future, we should
have a general solution to preserve the hole locations.

src/Juvix/Compiler/Internal/Builtins.hs

@@ -44,7 +44,7 @@ getBuiltinName ::
   Interval ->
   a ->
   Sem r Name
-getBuiltinName loc p = fromConcreteSymbol <$> getBuiltinSymbolHelper loc (toBuiltinPrim p)
+getBuiltinName loc p = fromConcreteSymbol loc <$> getBuiltinSymbolHelper loc (toBuiltinPrim p)
 
 checkBuiltinFunctionInfo ::
   forall r.

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

@@ -103,17 +103,17 @@ type ConstrName = Name
 
 type InductiveName = Name
 
-fromConcreteSymbol :: S.Symbol -> Name
-fromConcreteSymbol s = fromConcreteSymbolPretty (S.symbolText s) s
+fromConcreteSymbol :: Interval -> S.Symbol -> Name
+fromConcreteSymbol loc s = fromConcreteSymbolPretty loc (S.symbolText s) s
 
-fromConcreteSymbolPretty :: Text -> S.Symbol -> Name
-fromConcreteSymbolPretty pp s =
+fromConcreteSymbolPretty :: Interval -> Text -> S.Symbol -> Name
+fromConcreteSymbolPretty loc pp s =
   Name
     { _nameText = S.symbolText s,
       _nameId = s ^. S.nameId,
       _nameKind = getNameKind s,
       _nameKindPretty = getNameKindPretty s,
       _namePretty = pp,
-      _nameLoc = getLoc (s ^. S.nameConcrete),
+      _nameLoc = loc,
       _nameFixity = s ^. S.nameFixity
     }

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

@@ -0,0 +1,12 @@
+module Juvix.Compiler.Internal.Data.TypedIden where
+
+import Juvix.Compiler.Internal.Language
+import Juvix.Prelude
+
+data TypedIden = TypedIden
+  { _typedIden :: Iden,
+    _typedIdenType :: Expression
+  }
+  deriving stock (Data, Generic)
+
+makeLenses ''TypedIden

src/Juvix/Compiler/Internal/Extra.hs

@@ -14,6 +14,13 @@ import Juvix.Compiler.Internal.Language
 import Juvix.Compiler.Store.Internal.Data.InfoTable
 import Juvix.Prelude
 
+-- This is a hack to adjust location. It works only for identifiers. It should
+-- change the location of an arbitrary given expression to the given location.
+adjustLocation :: Interval -> Expression -> Expression
+adjustLocation loc = \case
+  ExpressionIden iden -> ExpressionIden (set (idenName . nameLoc) loc iden)
+  eh -> eh
+
 constructorArgTypes :: ConstructorInfo -> ([InductiveParameter], [FunctionParameter])
 constructorArgTypes i =
   ( i ^. constructorInfoInductiveParameters,
@@ -250,7 +257,9 @@ subsInstanceHoles s = umapM helper
   where
     helper :: Expression -> Sem r Expression
     helper le = case le of
-      ExpressionInstanceHole h -> clone (fromMaybe e (s ^. at h))
+      -- TODO: The location of the hole should be preserved
+      ExpressionInstanceHole h ->
+        adjustLocation (getLoc h) <$> clone (fromMaybe e (s ^. at h))
       _ -> return e
       where
         e = toExpression le
@@ -260,7 +269,9 @@ subsHoles s = umapM helper
   where
     helper :: Expression -> Sem r Expression
     helper le = case le of
-      ExpressionHole h -> clone (fromMaybe e (s ^. at h))
+      -- TODO: The location of the hole should be preserved
+      ExpressionHole h ->
+        adjustLocation (getLoc h) <$> clone (fromMaybe e (s ^. at h))
       _ -> return e
       where
         e = toExpression le

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

@@ -1,12 +1,14 @@
 module Juvix.Compiler.Internal.Extra.CoercionInfo
   ( module Juvix.Compiler.Store.Internal.Data.CoercionInfo,
     module Juvix.Compiler.Internal.Extra.CoercionInfo,
+    module Juvix.Compiler.Internal.Data.TypedIden,
   )
 where
 
 import Data.HashMap.Strict qualified as HashMap
 import Data.HashSet qualified as HashSet
 import Data.List qualified as List
+import Juvix.Compiler.Internal.Data.TypedIden
 import Juvix.Compiler.Internal.Extra.Base
 import Juvix.Compiler.Internal.Extra.InstanceInfo
 import Juvix.Compiler.Internal.Language
@@ -25,8 +27,8 @@ updateCoercionTable tab ci@CoercionInfo {..} =
 lookupCoercionTable :: CoercionTable -> Name -> Maybe [CoercionInfo]
 lookupCoercionTable tab name = HashMap.lookup name (tab ^. coercionTableMap)
 
-coercionFromTypedExpression :: TypedExpression -> Maybe CoercionInfo
-coercionFromTypedExpression TypedExpression {..}
+coercionFromTypedIden :: TypedIden -> Maybe CoercionInfo
+coercionFromTypedIden TypedIden {..}
   | null args = Nothing
   | otherwise = do
       tgt <- traitFromExpression metaVars (t ^. paramType)
@@ -36,11 +38,11 @@ coercionFromTypedExpression TypedExpression {..}
           { _coercionInfoInductive = _instanceAppHead,
             _coercionInfoParams = _instanceAppArgs,
             _coercionInfoTarget = tgt,
-            _coercionInfoResult = _typedExpression,
+            _coercionInfoResult = _typedIden,
             _coercionInfoArgs = args'
           }
   where
-    (args, e) = unfoldFunType _typedType
+    (args, e) = unfoldFunType _typedIdenType
     args' = init args
     t = List.last args
     metaVars = HashSet.fromList $ mapMaybe (^. paramName) args'

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

@@ -1,11 +1,13 @@
 module Juvix.Compiler.Internal.Extra.InstanceInfo
   ( module Juvix.Compiler.Store.Internal.Data.InstanceInfo,
     module Juvix.Compiler.Internal.Extra.InstanceInfo,
+    module Juvix.Compiler.Internal.Data.TypedIden,
   )
 where
 
 import Data.HashMap.Strict qualified as HashMap
 import Data.HashSet qualified as HashSet
+import Juvix.Compiler.Internal.Data.TypedIden
 import Juvix.Compiler.Internal.Extra.Base
 import Juvix.Compiler.Internal.Language
 import Juvix.Compiler.Store.Internal.Data.InstanceInfo
@@ -99,18 +101,18 @@ traitFromExpression metaVars e = case paramFromExpression metaVars e of
   Just (InstanceParamApp app) -> Just app
   _ -> Nothing
 
-instanceFromTypedExpression :: TypedExpression -> Maybe InstanceInfo
-instanceFromTypedExpression TypedExpression {..} = do
+instanceFromTypedIden :: TypedIden -> Maybe InstanceInfo
+instanceFromTypedIden TypedIden {..} = do
   InstanceApp {..} <- traitFromExpression metaVars e
   return $
     InstanceInfo
       { _instanceInfoInductive = _instanceAppHead,
         _instanceInfoParams = _instanceAppArgs,
-        _instanceInfoResult = _typedExpression,
+        _instanceInfoResult = _typedIden,
         _instanceInfoArgs = args
       }
   where
-    (args, e) = unfoldFunType _typedType
+    (args, e) = unfoldFunType _typedIdenType
     metaVars = HashSet.fromList $ mapMaybe (^. paramName) args
 
 checkNoMeta :: InstanceParam -> Bool

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

@@ -385,10 +385,10 @@ checkInstanceType ::
 checkInstanceType FunctionDef {..} = do
   ty <- strongNormalize _funDefType
   let mi =
-        instanceFromTypedExpression $
-          TypedExpression
-            { _typedType = ty ^. normalizedExpression,
-              _typedExpression = ExpressionIden (IdenFunction _funDefName)
+        instanceFromTypedIden $
+          TypedIden
+            { _typedIdenType = ty ^. normalizedExpression,
+              _typedIden = IdenFunction _funDefName
             }
   case mi of
     Just ii@InstanceInfo {..} -> do
@@ -435,10 +435,10 @@ checkCoercionType ::
 checkCoercionType FunctionDef {..} = do
   ty <- strongNormalize _funDefType
   let mi =
-        coercionFromTypedExpression
-          ( TypedExpression
-              { _typedType = ty ^. normalizedExpression,
-                _typedExpression = ExpressionIden (IdenFunction _funDefName)
+        coercionFromTypedIden
+          ( TypedIden
+              { _typedIdenType = ty ^. normalizedExpression,
+                _typedIden = IdenFunction _funDefName
               }
           )
   case mi of
@@ -1109,7 +1109,7 @@ inferLeftAppExpression mhint e = case e of
             typedLit litt blt ty = do
               from <- getBuiltinNameTypeChecker i blt
               ihole <- freshHoleImpl i ImplicitInstance
-              let ty' = fromMaybe ty mhint
+              let ty' = maybe ty (adjustLocation i) mhint
               inferExpression' (Just ty') $
                 foldApplication
                   (ExpressionIden (IdenFunction from))

src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/TypeChecking/Traits/Resolver.hs

@@ -41,7 +41,7 @@ resolveTraitInstance TypedHole {..} = do
   is <- lookupInstance ctab tab (ty ^. normalizedExpression)
   case is of
     [(cs, ii, subs)] ->
-      expandArity loc (subsIToE subs) (ii ^. instanceInfoArgs) (ii ^. instanceInfoResult)
+      expandArity' loc (subsIToE subs) (ii ^. instanceInfoArgs) (ii ^. instanceInfoResult)
         >>= applyCoercions loc cs
     [] ->
       throw (ErrNoInstance (NoInstance ty loc))
@@ -98,23 +98,23 @@ substitutionI subs p = case p of
     | otherwise ->
         return p
 
-instanceFromTypedExpression' :: InfoTable -> TypedExpression -> Maybe InstanceInfo
-instanceFromTypedExpression' tbl e = do
-  ii@InstanceInfo {..} <- instanceFromTypedExpression e
+instanceFromTypedIden' :: InfoTable -> TypedIden -> Maybe InstanceInfo
+instanceFromTypedIden' tbl e = do
+  ii@InstanceInfo {..} <- instanceFromTypedIden e
   guard (isTrait tbl _instanceInfoInductive)
   return ii
 
 varsToInstances :: InfoTable -> LocalVars -> [InstanceInfo]
 varsToInstances tbl LocalVars {..} =
   mapMaybe
-    (instanceFromTypedExpression' tbl . mkTyped)
+    (instanceFromTypedIden' tbl . mkTyped)
     (HashMap.toList _localTypes)
   where
-    mkTyped :: (VarName, Expression) -> TypedExpression
+    mkTyped :: (VarName, Expression) -> TypedIden
     mkTyped (v, ty) =
-      TypedExpression
-        { _typedType = ty,
-          _typedExpression = ExpressionIden (IdenVar v)
+      TypedIden
+        { _typedIdenType = ty,
+          _typedIden = IdenVar v
         }
 
 applyCoercions ::
@@ -133,10 +133,20 @@ applyCoercion ::
   Expression ->
   Sem r Expression
 applyCoercion loc (CoercionInfo {..}, subs) e = do
-  e' <- expandArity loc (subsIToE subs) _coercionInfoArgs _coercionInfoResult
+  e' <- expandArity' loc (subsIToE subs) _coercionInfoArgs _coercionInfoResult
   return $
     ExpressionApplication (Application e' e ImplicitInstance)
 
+expandArity' ::
+  (Members '[Error TypeCheckerError, NameIdGen] r) =>
+  Interval ->
+  Subs ->
+  [FunctionParameter] ->
+  Iden ->
+  Sem r Expression
+expandArity' loc subs params iden =
+  expandArity loc subs params (ExpressionIden (set (idenName . nameLoc) loc iden))
+
 expandArity ::
   (Members '[Error TypeCheckerError, NameIdGen] r) =>
   Interval ->

src/Juvix/Compiler/Store/Internal/Data/CoercionInfo.hs

@@ -10,7 +10,7 @@ data CoercionInfo = CoercionInfo
   { _coercionInfoInductive :: Name,
     _coercionInfoParams :: [InstanceParam],
     _coercionInfoTarget :: InstanceApp,
-    _coercionInfoResult :: Expression,
+    _coercionInfoResult :: Iden,
     _coercionInfoArgs :: [FunctionParameter]
   }
   deriving stock (Eq, Generic)

src/Juvix/Compiler/Store/Internal/Data/InstanceInfo.hs

@@ -44,7 +44,7 @@ instance NFData InstanceFun
 data InstanceInfo = InstanceInfo
   { _instanceInfoInductive :: InductiveName,
     _instanceInfoParams :: [InstanceParam],
-    _instanceInfoResult :: Expression,
+    _instanceInfoResult :: Iden,
     _instanceInfoArgs :: [FunctionParameter]
   }
   deriving stock (Eq, Generic)