More sensible backPropagation implementation

Author: MaximilianAlgehed

examples/Constrained/Examples/Basic.hs

@@ -343,3 +343,23 @@ manyInconsistentTrans = constrained' $ \ [var| a |] [var| b |] c d e [var| f |]
   , assert $ f >. 10
   , assert $ f <. b
   ]
+
+complicatedEither :: Specification (Either Int Int, (Either Int Int, Int, Int))
+complicatedEither = constrained' $ \ [var| i |] [var| t |] ->
+  [ caseOn i
+      (branch $ \ a -> a `elem_` lit [1..10])
+      (branch $ \ b -> b `elem_` lit [1..10])
+  , match t $ \ [var| k |] _ _ ->
+    [ k ==. i
+    , not_ $ k `elem_` lit [ Left j | j <- [1..9] ]
+    ]
+  ]
+
+pairCant :: Specification (Int, (Int, Int))
+pairCant = constrained' $ \ [var| i |] [var| p |] ->
+  [ assert $ i `elem_` lit [1..10]
+  , match p $ \ [var| k |] _ ->
+    [ k ==. i
+    , not_ $ k `elem_` lit [1..9]
+    ]
+  ]

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
@@ -833,17 +836,7 @@ mergeSolverStage (SolverStage x ps spec relevant) plan =
         normalizeSolverStage $ SolverStage
             y
             (ps ++ ps')
-            ( addToErrorSpec
-                ( NE.fromList
-                    ( [ "Solving var " ++ show x ++ " fails."
-                      , "Merging the Specs"
-                      , "   1. " ++ show spec
-                      , "   2. " ++ show spec'
-                      ]
-                    )
-                )
-                (spec <> spec')
-            )
+            (spec <> spec')
             (relevant <> relevant')
       Nothing -> stage
   | stage@(SolverStage y ps' spec' relevant') <- plan
@@ -897,27 +890,21 @@ backPropagation relevant (SolverPlan initplan) = SolverPlan (go [] (reverse init
     go acc [] = acc
     go acc (s@(SolverStage (x :: Var a) ps spec _) : plan) = go (s : acc) plan'
       where
-        newStages = concatMap (newStage spec) ps
+        newStages = concatMap newStage ps
         plan' = foldr mergeSolverStage plan newStages
+
         -- Note use of the Term Pattern Equal
-        newStage specl (Assert (Equal (V x') t)) =
-          termVarEqCases specl x' t
-        newStage specr (Assert (Equal t (V x'))) =
-          termVarEqCases specr x' t
-        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
-          | Just Refl <- eqVar x x'
-          , [Name y] <- Set.toList $ freeVarSet t
-          , Result ctx <- toCtx y t =
-              [SolverStage y [] (propagateSpec specx ctx)
-               (Set.insert (Name x') relevant)]
-        termVarEqCases _ _ _ = []
+        newStage (Assert (Equal tl tr))
+          | [Name xl] <- Set.toList $ freeVarSet tl
+          , [Name xr] <- Set.toList $ freeVarSet tr
+          , Name x `elem` [Name xl, Name xr]
+          , Result ctxL <- toCtx xl tl
+          , Result ctxR <- toCtx xr tr
+          = case (eqVar x xl, eqVar x xr) of
+              (Just Refl, _) -> [SolverStage xr [] (propagateSpec (forwardPropagateSpec spec ctxL) ctxR) (Set.insert (Name x) relevant)]
+              (_, Just Refl) -> [SolverStage xl [] (propagateSpec (forwardPropagateSpec spec ctxR) ctxL) (Set.insert (Name x) relevant)]
+              _ -> error "The impossible happened"
+        newStage _ = []
 
 -- | Function symbols for `(==.)`
 data EqW :: [Type] -> Type -> Type where
@@ -1329,3 +1316,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
 -- ================================================================

test/Constrained/Tests.hs

@@ -59,6 +59,8 @@ testAll = hspec $ tests False
 tests :: Bool -> Spec
 tests nightly =
   describe "constrained" . modifyMaxSuccess (\ms -> if nightly then ms * 10 else ms) $ do
+    testSpec "complicatedEither" complicatedEither
+    testSpec "pairCatn" pairCant
     -- TODO: double-shrinking
     testSpecNoShrink "reifiesMultiple" reifiesMultiple
     testSpec "assertReal" assertReal