@@ -77,20 +77,40 @@ runConstraintEnv ::
7777 Members '[State VariableMap, State (MKM.Map Pred)] effs =>
7878 Eff (SubstituteEnv : Writer Substitution : Writer Constraint : State (Seq Constraint) : effs) a ->
7979 Eff effs ((a, Substitution), Seq Constraint)
80−runConstraintEnv =
81− runState @(Seq Constraint) Empty
82− . interpret @(Writer Constraint) (\(Tell x) -> modify (x :<|))
83− . runWriter @Substitution
84− . interpret @SubstituteEnv
85− ( \(SubstituteEnv x) -> do
86− modify @(Seq Constraint) (sub x <$>)
87− -- TODO: I have a feeling that variables in the env should already be closed over.
88− -- (aka this is redundant)
89− modify @VariableMap (sub x <$>)
90− modify @(MKM.Map Pred) (sub x)
91− tell x
92− return ()
93− )
80+runConstraintEnv f = do
81+ ((x, subs), cs) <-
82+ runState @(Seq Constraint) Empty
83+ . interpret @(Writer Constraint) (\(Tell x) -> modify (x :<|))
84+ . runWriter @Substitution
85+ . interpret @SubstituteEnv
86+ ( \(SubstituteEnv x) -> do
87+ -- modify @(Seq Constraint) (sub x <$>)
88+ -- TODone: I have a feeling that variables in the env should already be closed over.
89+ -- (aka this is redundant)
90+ --
91+ -- No, they can have type variables. E g when they are arguments in a function.
92+ modify @VariableMap (sub x <$>)
93+ modify @(MKM.Map Pred) (sub x)
94+ tell x
95+ return ()
96+ )
97+ $ f
98+ return ((x, subs), sub subs <$> cs)
99+
100+captureSubstitutions ::
101+ Members '[SubstituteEnv] effs =>
102+ Eff (Writer Substitution : effs) a ->
103+ Eff effs (a, Substitution)
104+captureSubstitutions f = do
105+ (x, s) <-
106+ runWriter @Substitution
107+ . interpose @SubstituteEnv
108+ ( \(SubstituteEnv x) -> do
109+ tell x
110+ subEnv x
111+ )
112+ $ f
113+ return (x, s)
94114
95115data SubstituteEnv a where
96116 SubstituteEnv :: Substitution -> SubstituteEnv ()
@@ -166,7 +186,7 @@ prettySubstitution = A.toJSON . M.mapKeysMonotonic (TL.pack . show) . fmap rende
166186
167187-- | @older <> newer@
168188instance Semigroup Substitution where
169− x'@(Substitution x) <> (Substitution y) = Substitution (fmap (sub x') y <> x)
189+ x'@(Substitution x) <> y'@(Substitution y) = Substitution (M.unionWithKey (error "invariant") (fmap (sub y') x) (fmap (sub x') y))
170190
171191instance Monoid Substitution where
172192 mempty = Substitution mempty
@@ -598,8 +618,12 @@ inferComponents n ((S.toList -> paths) : otherComponents) = do
598618 return $ fmap fst <$> tvsTypes
599619 -- TODO: There should probably be some smarter predicate filtering
600620 schemes <- types & traverse . traverse %%~ (\t -> solveAndClose p $ sub s $ preds :=> t)
601− fmap (M.unionWith mergeBindings schemes) . withBindings (toScheme' <$> schemes) $
602− inferComponents n otherComponents
621+
622+ (otherBindings, s') <-
623+ captureSubstitutions $
624+ withBindings (toScheme' <$> schemes) $
625+ inferComponents n otherComponents
626+ return $ M.unionWith mergeBindings (fmap (sub s') <$> schemes) otherBindings
603627 where
604628 bindingToNExpr :: NonEmpty VarName -> NExprLoc -> Binding NExpr
605629 bindingToNExpr name expr = NamedVar (StaticKey <$> name) (stripAnnotation expr) nullPos
@@ -724,6 +748,7 @@ inferGeneral' toPretty f = do
724748 traceValue "substitutions_after_solving" $ prettySubstitution s'
725749 let s = subs <> s'
726750 returnedS = Substitution . M.filterWithKey (\k _ -> k `RS.notMember` range) . unSubstitution $ s
751+ traceValue "composed_substitution" $ prettySubstitution s
727752 traceSubtree "returned" $ do
728753 traceValue "substitutions" $ prettySubstitution returnedS
729754 -- Filter for optimization purposes