@@ -24,6 +24,7 @@ import Control.Monad.Freer.TreeTracer
2424import Control.Monad.Freer.Writer
2525import Data.Act
2626import qualified Data.Aeson as A
27+import Data.Either
2728import Data.Fix
2829import Data.Foldable
2930import Data.Group hiding ((~~))
@@ -61,7 +62,7 @@ getType r =
6162 . runReader nullSpan
6263 . evalState @(MKM.Map Pred) mempty
6364 . interpret @(Writer Pred) (\(Tell x) -> modify (<> MKM.singleton x))
64− . evalState (M.empty :: VariableMap)
65+ . runReader (M.empty :: VariableMap)
6566 . evalState @Substitution mempty
6667 . runState @[(TypeVariable, SrcSpan)] []
6768 . interpret @(Writer (TypeVariable, SrcSpan)) (\(Tell x) -> modify (x :))
@@ -86,6 +87,11 @@ runFreshSrc =
8687 src <- ask @SrcSpan
8788 tell (v, src)
8889 return v
90+ FreshSub x -> do
91+ v <- send $ FreshSub x
92+ src <- ask @SrcSpan
93+ tell (v, src)
94+ return v
8995 MinFreshFromNow -> send MinFreshFromNow
9096 )
9197
@@ -183,7 +189,7 @@ v <<- x = modify (Substitution (M.singleton v x) <>)
183189unifyWithPriority ::
184190 -- | The range of type variables to interpret as being lower
185191 -- priority (try to replace them if possible)
186− RS.RSet TypeVariable ->
192+ RS.RSet RangeTypeVariable ->
187193 Constraint ->
188194 UnifyM ()
189195unifyWithPriority range (lhs :~ rhs) = do
@@ -191,7 +197,7 @@ unifyWithPriority range (lhs :~ rhs) = do
191197 case sub s (lhs, rhs) of
192198 (x, y) | x == y -> return mempty
193199 (NTypeVariable x, NTypeVariable y) -> do
194− if y `RS.member` range
200+ if rangeTypeVariable y `RS.member` range
195201 then y <<- NTypeVariable x
196202 else x <<- NTypeVariable y
197203 (NTypeVariable t, x) -> t <<- x
@@ -205,13 +211,13 @@ unifyWithPriority range (lhs :~ rhs) = do
205211 unifyWithPriority range $ y1 :~ y2
206212 (x, y) -> tell (CanNotUnify x y)
207213
208−unifySchemes :: RS.RSet TypeVariable -> [(Scheme, Scheme)] -> UnifyM ()
214+unifySchemes :: RS.RSet RangeTypeVariable -> [(Scheme, Scheme)] -> UnifyM ()
209215unifySchemes _ [] = pure mempty
210216unifySchemes range ((x, y) : rest) = do
211217 unifyScheme range x y
212218 unifySchemes range rest
213219
214−unifyScheme :: RS.RSet TypeVariable -> Scheme -> Scheme -> UnifyM ()
220+unifyScheme :: RS.RSet RangeTypeVariable -> Scheme -> Scheme -> UnifyM ()
215221unifyScheme range (xs :=> x) (ys :=> y) = do
216222 solveWithPriority range (pure (x :~ y))
217223 s <- get
@@ -229,7 +235,7 @@ solve = solveWithPriority mempty
229235solveWithPriority ::
230236 -- | The range of type variables to interpret as being lower
231237 -- priority (try to replace them if possible)
232− RS.RSet TypeVariable ->
238+ RS.RSet RangeTypeVariable ->
233239 Seq Constraint ->
234240 UnifyM ()
235241solveWithPriority _ Empty = return mempty
@@ -267,7 +273,7 @@ type InferEffs =
267273 Writer PredicateError,
268274 Writer Pred,
269275 State (MKM.Map Pred),
270− State VariableMap,
276+ Reader VariableMap,
271277 Writer UnifyingError,
272278 Writer InferError,
273279 State Substitution,
@@ -304,17 +310,12 @@ type VariableMap = Map VarName Scheme
304310(~~) :: Member (Writer Constraint) r => NType -> NType -> Eff r ()
305311lhs ~~ rhs = tell $ lhs :~ rhs
306312
307−-- TODO: shadowed bindings are not properly restored.
308313withBindings ::
309− (Member (State VariableMap) effs) =>
314+ (Member (Reader VariableMap) effs) =>
310315 VariableMap ->
311316 Eff effs a ->
312317 Eff effs a
313−withBindings bindings m = do
314− modify (bindings <>)
315− a <- m
316− modify @VariableMap (`M.difference` bindings)
317− return a
318+withBindings bindings = local (bindings <>)
318319
319320instantiate ::
320321 forall effs.
@@ -356,7 +357,7 @@ infer (Fix (Compose (Ann src x))) = local (const src) $ case x of
356357 NEnvPath {} -> return . NAtomic $ Path
357358 NStr {} -> return $ NAtomic String
358359 NSym var ->
359− get @VariableMap
360+ ask @VariableMap
360361 >>= maybe (scheme . NTypeVariable <$> throwTV (UndefinedVariable var)) return . M.lookup var
361362 >>= instantiate -- I have no idea if this covers all cases.
362363 NList xs -> do
@@ -413,8 +414,7 @@ infer (Fix (Compose (Ann src x))) = local (const src) $ case x of
413414 resT <-
414415 NTypeVariable <$> case path of
415416 Just p -> do
416− (set', var) <- lookupAttrSet p
417− setT ~~ set'
417+ (_, var) <- lookupAttrSet setT p
418418 return var
419419 Nothing -> fresh
420420 case def of
@@ -563,27 +563,30 @@ inferComponents :: NormalizedBinding NExprLoc -> [Set (NonEmpty VarName)] -> Inf
563563inferComponents _ [] = pure M.empty
564564inferComponents n ((S.toList -> paths) : otherComponents) = do
565565 let bindings = catMaybes $ paths <&> \path -> (path,) <$> lookupExpr path n
566− (types, p, preds, s) <-
567− traceSubtree (showExpr . bindingsToNExpr $ bindings)
568− . inferGeneral' (NAttrSet . fmap toScheme)
569− $ do
570− tvsExprs <- for bindings (\(a, b) -> (a,b,) . NTypeVariable <$> fresh)
571− let tvsExprs' =
572− foldl' (M.unionWith mergeBindings) M.empty $
573− fmap (\(a, b, c) -> normalize a (b, c)) tvsExprs
574− tvs = toScheme . fmap snd <$> tvsExprs'
575−
576− -- (type, tv)
577− tvsTypes <-
578− withBindings tvs $
579− tvsExprs' & traverse . traverse . _1 %%~ infer
580− for_ tvsTypes $ traverse (uncurry (~~))
581− return $ fmap fst <$> tvsTypes
582− -- TODO: There should probably be some smarter predicate filtering
583− schemes <- types & traverse . traverse %%~ (\t -> solveAndClose p $ sub s $ preds :=> t)
566+ (schemes, schemes') <- traceSubtree (showExpr . bindingsToNExpr $ bindings) $ do
567+ (types, p, preds, s) <-
568+ inferGeneral' (NAttrSet . fmap toScheme) $
569+ do
570+ tvsExprs <- for bindings (\(a, b) -> (a,b,) . NTypeVariable <$> fresh)
571+ let tvsExprs' =
572+ foldl' (M.unionWith mergeBindings) M.empty $
573+ fmap (\(a, b, c) -> normalize a (b, c)) tvsExprs
574+ tvs = toScheme . fmap snd <$> tvsExprs'
575+
576+ -- (type, tv)
577+ tvsTypes <-
578+ withBindings tvs $
579+ tvsExprs' & traverse . traverse . _1 %%~ infer
580+ for_ tvsTypes $ traverse (uncurry (~~))
581+ return $ fmap fst <$> tvsTypes
582+ -- TODO: There should probably be some smarter predicate filtering
583+ schemes <- types & traverse . traverse %%~ (\t -> solveAndClose p $ sub s $ preds :=> t)
584+ let schemes' = toScheme' <$> schemes
585+ traceValue "resulting_type" $ renderPretty . scheme $ NAttrSet schemes'
586+ return (schemes, schemes')
584587
585588 otherBindings <-
586− withBindings (toScheme' <$> schemes) $ inferComponents n otherComponents
589+ withBindings schemes' $ inferComponents n otherComponents
587590 return $ M.unionWith mergeBindings schemes otherBindings
588591 where
589592 bindingToNExpr :: NonEmpty VarName -> NExprLoc -> Binding NExpr
@@ -690,13 +693,15 @@ inferGeneral :: NExprLoc -> InferM Scheme
690693inferGeneral x@(Fix (Compose (Ann src _))) =
691694 local (const src) . traceSubtree (showExpr . stripAnnotation $ x) $ do
692695 (t, p, keep, s) <- inferGeneral' id $ infer x
693− solveAndClose p $ sub s $ keep :=> t
696+ res <- solveAndClose p $ sub s $ keep :=> t
697+ traceValue "resuting_type" $ renderPretty res
698+ return res
694699
695700inferGeneral' ::
696701 (Members InferEffs r, Pretty y) =>
697702 (x -> y) ->
698703 Eff (Writer Constraint : r) x ->
699− Eff r (x, RS.RSet TypeVariable, [Pred], Substitution)
704+ Eff r (x, RS.RSet RangeTypeVariable, [Pred], Substitution)
700705inferGeneral' toPretty f = do
701706 ((x, cs), range) <-
702707 traceSubtree "subexpressions" . registerTypeVariables . runSeqWriter @(Seq Constraint) $ f
@@ -710,7 +715,7 @@ inferGeneral' toPretty f = do
710715 get >>= traceValue "substitutions_after_solving" . prettySubstitution
711716 (s, returnedS) <-
712717 get @Substitution
713− <&> join bimap Substitution . M.partitionWithKey (\k _ -> k `RS.member` range) . unSubstitution
718+ <&> join bimap Substitution . M.partitionWithKey (\k _ -> rangeTypeVariable k `RS.member` range) . unSubstitution
714719 put returnedS
715720 -- modify @Substitution $ Substitution . M.filterWithKey (\k _ -> k `RS.notMember` range) . unSubstitution
716721 -- returnedS <- get @Substitution
@@ -718,9 +723,8 @@ inferGeneral' toPretty f = do
718723 traceValue "substitutions" $ prettySubstitution returnedS
719724 -- Filter for optimization purposes
720725 preds <- get @(MKM.Map Pred)
721− traceValue "subbed_preds" $ showPreds preds
722726 let newRange = range RS.\\ rangeFromSet (free returnedS)
723− (keep', toss) = MKM.partition (`RS.member` newRange) preds
727+ (keep', toss) = MKM.partition ((`RS.member` newRange) . rangeTypeVariable) preds
724728 keep = MKM.toList keep'
725729 -- At this point the only predicates only contain type variables that are only used
726730 -- in the expression. This means that all of the generated constraints and
@@ -733,8 +737,8 @@ inferGeneral' toPretty f = do
733737 showPreds :: MKM.Map Pred -> Map String String
734738 showPreds = M.mapKeys show . fmap show . MKM.toMap
735739
736−rangeFromSet :: (Ord x, Enum x) => S.Set x -> RS.RSet x
737−rangeFromSet = RS.fromAscList . S.toAscList
740+rangeFromSet :: S.Set TypeVariable -> RS.RSet RangeTypeVariable
741+rangeFromSet = RS.fromAscList . fmap rangeTypeVariable . S.toAscList
738742
739743solveAndClose ::
740744 Members
@@ -743,20 +747,21 @@ solveAndClose ::
743747 Writer UnifyingError
744748 ]
745749 eff =>
746− RS.RSet TypeVariable ->
750+ RS.RSet RangeTypeVariable ->
747751 Scheme ->
748752 Eff eff Scheme
749−solveAndClose range (cs :=> t) = do
753+solveAndClose range (cs :=> t) = evalState @Substitution mempty $ do
750754 (((solvedPreds, range'), constraints), collectedPreds) <-
751755 runSeqWriter @[Pred]
752756 . runSeqWriter @(Seq Constraint)
753757 . registerTypeVariables
754− $ solveConstraints cs
758+ $ solveConstraints range cs
755759 let newRange = range `RS.union` range'
756760 case (constraints, collectedPreds) of
757761 (Empty, []) -> return . close newRange $ solvedPreds :=> t
758762 _ -> do
759− ((), s) <- runState @Substitution mempty $ solveWithPriority newRange constraints
763+ solveWithPriority newRange constraints
764+ s <- get
760765 solveAndClose newRange $ sub s $ (collectedPreds <> solvedPreds) :=> t
761766
762767solveConstraints ::
@@ -764,30 +769,43 @@ solveConstraints ::
764769 '[ Fresh,
765770 Writer Pred,
766771 Writer PredicateError,
767− Writer Constraint
772+ Writer Constraint,
773+ Writer UnifyingError,
774+ State Substitution
768775 ]
769776 effs =>
777+ RS.RSet RangeTypeVariable ->
770778 [Pred] ->
771779 Eff effs [Pred]
772−solveConstraints ps = do
773− flip Control.Monad.filterM ps $ \case
774− HasField (NAttrSet ts) (f, t) -> case M.lookup f ts of
775− Nothing -> do
776− tell $ KeyNotPresent f (NAttrSet ts)
777− return False
778− Just t' -> do
779− x <- instantiate t
780− y <- instantiate t'
781− x ~~ y
782− return False
783− HasField (NTypeVariable _) _ -> return True
784− HasField t _ -> do
785− tell $ NotAnAttributeSet t
786− return False
787− (Update (NAttrSet a) (NAttrSet b) t) -> do
788− (NAttrSet $ b <> a) ~~ t
789− return False
790− _ -> return True
780+solveConstraints _ [] = pure []
781+solveConstraints range (p : ps) = case p of
782+ HasField (NAttrSet ts) (f, t) -> case M.lookup f ts of
783+ Nothing -> do
784+ tell $ KeyNotPresent f (NAttrSet ts)
785+ solveConstraints range ps
786+ Just t' -> do
787+ x <- instantiate t
788+ y <- instantiate t'
789+ x ~~ y
790+ solveConstraints range ps
791+ HasField tv@(NTypeVariable _) (k, v) -> do
792+ let sameAttrSet :: Pred -> Either Scheme Pred
793+ sameAttrSet = \case
794+ HasField t (k', v') | t == tv, k == k' -> Left v'
795+ x -> Right x
796+ (same, different) = partitionWith sameAttrSet ps
797+ (unifyScheme range v) `traverse_` same
798+ (p :) <$> solveConstraints range different
799+ HasField t _ -> do
800+ tell $ NotAnAttributeSet t
801+ solveConstraints range ps
802+ (Update (NAttrSet a) (NAttrSet b) t) -> do
803+ (NAttrSet $ b <> a) ~~ t
804+ solveConstraints range ps
805+ _ -> (p :) <$> solveConstraints range ps
806+
807+partitionWith :: (a -> Either b c) -> [a] -> ([b], [c])
808+partitionWith f = partitionEithers . map f
791809
792810newtype DeBruijnContext = DeBruijnContext (Sum Int)
793811 deriving newtype (Eq, Ord, Show, Semigroup, Monoid, Group, Num)
@@ -797,7 +815,7 @@ instance Act DeBruijnContext DeBruijn where
797815
798816type TDeBruijnMap = ShiftedMap DeBruijnContext TypeVariable DeBruijn
799817
800−close :: TraversableNTypes x => RS.RSet TypeVariable -> x -> x
818+close :: TraversableNTypes x => RS.RSet RangeTypeVariable -> x -> x
801819close range =
802820 run
803821 . evalState @TDeBruijnMap mempty
@@ -811,7 +829,7 @@ close range =
811829 traverseNTypesWith bndCtx %%~ \case
812830 x@(TBruijn _) -> return x
813831 x@(TAtomic _) -> return x
814− (TTypeVariable tv) | tv `RS.member` range -> TBruijn <$> newbruijn tv
832+ (TTypeVariable tv) | rangeTypeVariable tv `RS.member` range -> TBruijn <$> newbruijn tv
815833 x@(TTypeVariable _) -> return x
816834
817835 bndCtx m = do
@@ -866,18 +884,22 @@ lookupAttrSet ::
866884 Fresh
867885 ]
868886 r =>
887+ NType ->
869888 AttrSetPath ->
870889 Eff r (NType, TypeVariable)
871−lookupAttrSet (NESingle p) = do
872− set <- NTypeVariable <$> fresh
873− var <- fresh
874− tell (HasField set (p, scheme . NTypeVariable $ var))
875− return (set, var)
876−lookupAttrSet (p :|| ps) = do
877− set <- NTypeVariable <$> fresh
878− (innerSet, var) <- lookupAttrSet ps
879− tell $ HasField set (p, scheme innerSet)
880− return (set, var)
890+lookupAttrSet n = \case
891+ (NESingle p) -> do
892+ var <- getSubTV n
893+ tell (HasField n (p, scheme . NTypeVariable $ var))
894+ return (n, var)
895+ (p :|| ps) -> do
896+ (innerSet, var) <- lookupAttrSet n ps
897+ set <- NTypeVariable <$> getSubTV n
898+ tell $ HasField set (p, scheme innerSet)
899+ return (set, var)
900+ where
901+ getSubTV (NTypeVariable tv) = freshSub tv
902+ getSubTV _ = fresh
881903
882904-- NonEmpty patterns
883905