More sensible backPropagation implementation

Author: MaximilianAlgehed

src/Constrained/Base.hs

@@ -32,10 +32,11 @@ module Constrained.Base (
   pattern (:<:),
   pattern (:>:),
   pattern Unary,
-  Ctx,
+  Ctx(..),
   toCtx,
   flipCtx,
   fromListCtx,
+  ctxHasSpec,
 
   -- * Useful function symbols and patterns for building custom rewrite rules
   fromGeneric_,

src/Constrained/Generation.hs

@@ -49,6 +49,9 @@ module Constrained.Generation (
   EqW (..),
   SumSpec (..),
   pattern SumSpec,
+
+  mapSpec,
+  forwardPropagateSpec,
 ) where
 
 import Constrained.AbstractSyntax
@@ -907,10 +910,10 @@ backPropagation relevant (SolverPlan initplan) = SolverPlan (go [] (reverse init
         newStage _ _ = []
 
         termVarEqCases :: HasSpec b => Specification a -> Var b -> Term b -> [SolverStage]
-        termVarEqCases (MemberSpec vs) x' t
-          | Set.singleton (Name x) == freeVarSet t =
-              [SolverStage x' [] (MemberSpec (NE.nub (fmap (\v -> errorGE $ runTerm (Env.singleton x v) t) vs)))
-               (Set.insert (Name x') relevant)]
+        termVarEqCases specx x' t
+          | Set.singleton (Name x) == freeVarSet t
+          , Result ctx <- toCtx x t =
+            [SolverStage x' [] (forwardPropagateSpec specx ctx) (Set.insert (Name x) relevant)]
         termVarEqCases specx x' t
           | Just Refl <- eqVar x x'
           , [Name y] <- Set.toList $ freeVarSet t
@@ -1329,3 +1332,34 @@ fromGESpec ge = case ge of
   Result s -> s
   GenError xs -> ErrorSpec (catMessageList xs)
   FatalError es -> error $ catMessages es
+
+-- TODO: move this somewhere sensible
+
+-- | Functor like property for Specification, but instead of a Haskell function (a -> b),
+--   it takes a function symbol (t '[a] b) from a to b.
+--   Note, in this context, a function symbol is some constructor of a witnesstype.
+--   Eg. ProdFstW, InjRightW, SingletonW, etc. NOT the lifted versions like fst_ singleton_,
+--   which construct Terms. We had to wait until here to define this because it
+--   depends on Semigroup property of Specification, and Asserting equality
+mapSpec ::
+  forall t a b.
+  AppRequires t '[a] b =>
+  t '[a] b ->
+  Specification a ->
+  Specification b
+mapSpec f (ExplainSpec es s) = explainSpec es (mapSpec f s)
+mapSpec f TrueSpec = mapTypeSpec f (emptySpec @a)
+mapSpec _ (ErrorSpec err) = ErrorSpec err
+mapSpec f (MemberSpec as) = MemberSpec $ NE.nub $ fmap (semantics f) as
+mapSpec f (SuspendedSpec x p) =
+  constrained $ \x' ->
+    Exists (\_ -> fatalError "mapSpec") (x :-> fold [Assert $ (x' ==. appTerm f (V x)), p])
+mapSpec f (TypeSpec ts cant) = mapTypeSpec f ts <> notMemberSpec (map (semantics f) cant)
+
+-- TODO generalizeme!
+forwardPropagateSpec :: HasSpec a => Specification a -> Ctx a b -> Specification b
+forwardPropagateSpec s CtxHOLE = s
+forwardPropagateSpec s (CtxApp f (c :? Nil))
+  | Evidence <- ctxHasSpec c = mapSpec f (forwardPropagateSpec s c)
+forwardPropagateSpec _ _ = TrueSpec
+

src/Constrained/TheKnot.hs

@@ -51,7 +51,6 @@ module Constrained.TheKnot (
   rangeSize,
   hasSize,
   genInverse,
-  mapSpec,
   between,
 
   -- * Patterns
@@ -77,7 +76,6 @@ import Constrained.SumList
 -- Because it is mutually recursive with something else in here.
 import Constrained.Syntax
 import Control.Applicative
-import Control.Monad
 import Data.Foldable
 import Data.Kind
 import Data.List (nub)
@@ -99,27 +97,6 @@ ifElse b p q = whenTrue b p <> whenTrue (not_ b) q
 
 -- =======================================================================================
 
--- | Functor like property for Specification, but instead of a Haskell function (a -> b),
---   it takes a function symbol (t '[a] b) from a to b.
---   Note, in this context, a function symbol is some constructor of a witnesstype.
---   Eg. ProdFstW, InjRightW, SingletonW, etc. NOT the lifted versions like fst_ singleton_,
---   which construct Terms. We had to wait until here to define this because it
---   depends on Semigroup property of Specification, and Asserting equality
-mapSpec ::
-  forall t a b.
-  AppRequires t '[a] b =>
-  t '[a] b ->
-  Specification a ->
-  Specification b
-mapSpec f (ExplainSpec es s) = explainSpec es (mapSpec f s)
-mapSpec f TrueSpec = mapTypeSpec f (emptySpec @a)
-mapSpec _ (ErrorSpec err) = ErrorSpec err
-mapSpec f (MemberSpec as) = MemberSpec $ NE.nub $ fmap (semantics f) as
-mapSpec f (SuspendedSpec x p) =
-  constrained $ \x' ->
-    Exists (\_ -> fatalError "mapSpec") (x :-> fold [Assert $ (x' ==. appTerm f (V x)), p])
-mapSpec f (TypeSpec ts cant) = mapTypeSpec f ts <> notMemberSpec (map (semantics f) cant)
-
 -- ================================================================
 -- HasSpec for Products
 -- ================================================================