Resurrect shrinking and improve it (#51)

Author: MaximilianAlgehed

examples/Constrained/Examples/BinTree.hs

@@ -76,6 +76,8 @@ instance HasSpec a => HasSpec (BinTree a) where
       : (BinNode left a <$> shrinkWithTypeSpec s right)
       ++ ((\l -> BinNode l a right) <$> shrinkWithTypeSpec s left)
 
+  fixupWithTypeSpec _ _ = Nothing
+
   cardinalTypeSpec _ = mempty
 
   toPreds t (BinTreeSpec msz s) =

src/Constrained/Base.hs

@@ -488,6 +488,12 @@ class
   -- | Shrink an `a` with the aide of a `TypeSpec`
   shrinkWithTypeSpec :: TypeSpec a -> a -> [a]
 
+  -- | Try to make an `a` conform to `TypeSpec` with minimal changes. When
+  -- `fixupWithSpec ts a` returns `Just a'`, it should be the case that
+  -- `conformsTo a' ts`. There are no constraints in the `Nothing` case. A
+  -- non-trivial implementation of this function is important for shrinking.
+  fixupWithTypeSpec :: TypeSpec a -> a -> Maybe a
+
   -- | Convert a spec to predicates:
   -- The key property here is:
   --   ∀ a. a `conformsTo` spec == a `conformsTo` constrained (\t -> toPreds t spec)
@@ -598,6 +604,13 @@ class
     [a]
   shrinkWithTypeSpec spec a = map fromSimpleRep $ shrinkWithTypeSpec spec (toSimpleRep a)
 
+  default fixupWithTypeSpec ::
+    GenericallyInstantiated a =>
+    TypeSpec a ->
+    a ->
+    Maybe a
+  fixupWithTypeSpec spec a = fromSimpleRep <$> fixupWithTypeSpec spec (toSimpleRep a)
+
   default cardinalTypeSpec ::
     GenericallyInstantiated a =>
     TypeSpec a ->
@@ -618,6 +631,7 @@ instance HasSpec Bool where
   cardinalTypeSpec _ = equalSpec 2
   cardinalTrueSpec = equalSpec 2
   shrinkWithTypeSpec _ = shrink
+  fixupWithTypeSpec _ = Just
   conformsTo _ _ = True
   toPreds _ _ = TruePred
 
@@ -627,6 +641,7 @@ instance HasSpec () where
   combineSpec _ _ = typeSpec ()
   _ `conformsTo` _ = True
   shrinkWithTypeSpec _ _ = []
+  fixupWithTypeSpec _ _ = pure ()
   genFromTypeSpec _ = pure ()
   toPreds _ _ = TruePred
   cardinalTypeSpec _ = MemberSpec (pure 1)

src/Constrained/Conformance.hs

@@ -14,6 +14,7 @@ module Constrained.Conformance (
   conformsToSpec,
   conformsToSpecE,
   satisfies,
+  checkPredE,
   checkPredsE,
 ) where
 

src/Constrained/Generation.hs

@@ -25,6 +25,7 @@ module Constrained.Generation (
   genFromSpecT,
   genFromSpecWithSeed,
   shrinkWithSpec,
+  fixupWithSpec,
   simplifySpec,
 
   -- ** Debugging
@@ -155,15 +156,71 @@ shrinkWithSpec (simplifySpec -> spec) a = filter (`conformsToSpec` spec) $ case
   ExplainSpec _ s -> shrinkWithSpec s a
   -- TODO: filter on can't if we have a known to be sound shrinker
   TypeSpec s _ -> shrinkWithTypeSpec s a
-  -- TODO: The better way of doing this is to compute the dependency graph,
-  -- shrink one variable at a time, and fixup the rest of the variables
-  SuspendedSpec {} -> shr a
+  SuspendedSpec x p -> shrinkFromPreds p x a ++ shr a
   MemberSpec {} -> shr a
   TrueSpec -> shr a
   ErrorSpec {} -> []
   where
     shr = shrinkWithTypeSpec (emptySpec @a)
 
+shrinkFromPreds :: HasSpec a => Pred -> Var a -> a -> [a]
+shrinkFromPreds p
+  | Result plan <- prepareLinearization p = \x a -> listFromGE $ do
+      -- NOTE: we do this to e.g. guard against bad construction functions in Exists
+      case checkPredE (Env.singleton x a) (NE.fromList []) p of
+        Nothing -> pure ()
+        Just err -> explainNE err $ fatalError "Trying to shrink a bad value, don't do that!"
+      -- Get an `env` for the original value
+      initialEnv <- envFromPred (Env.singleton x a) p
+      return
+        [ a'
+        | -- Shrink the initialEnv
+        env' <- shrinkEnvFromPlan initialEnv plan
+        , -- Get the value of the constrained variable `x` in the shrunk env
+        Just a' <- [Env.lookup env' x]
+        , -- NOTE: this is necessary because it's possible that changing
+        -- a particular value in the env during shrinking might not result
+        -- in the value of `x` changing and there is no better way to know than
+        -- to do this.
+        a' /= a
+        ]
+  | otherwise = error "Bad pred"
+
+-- Start with a valid Env for the plan and try to shrink it
+shrinkEnvFromPlan :: Env -> SolverPlan -> [Env]
+shrinkEnvFromPlan initialEnv SolverPlan {..} = go mempty solverPlan
+  where
+    go :: Env -> [SolverStage] -> [Env]
+    go _ [] = [] -- In this case we decided to keep every variable the same so nothing to return
+    go env ((unsafeSubstStage env -> SolverStage {..}) : plan) = do
+      Just a <- [Env.lookup initialEnv stageVar]
+      -- Two cases:
+      --  - either we shrink this value and try to fixup every value later on in the plan or
+      [ fixedEnv
+        | a' <- shrinkWithSpec stageSpec a
+        , let env' = Env.extend stageVar a' env
+        , Just fixedEnv <- [fixupPlan env' plan]
+        ]
+        --  - we keep this value the way it is and try to shrink some later value
+        ++ go (Env.extend stageVar a env) plan
+
+    -- Fix the rest of the plan given an environment `env` for the plan so far
+    fixupPlan :: Env -> [SolverStage] -> Maybe Env
+    fixupPlan env [] = pure env
+    fixupPlan env ((unsafeSubstStage env -> SolverStage {..}) : plan) =
+      case Env.lookup (env <> initialEnv) stageVar >>= fixupWithSpec stageSpec of
+        Nothing -> Nothing
+        Just a -> fixupPlan (Env.extend stageVar a env) plan
+
+-- Try to fix a value w.r.t a specification
+fixupWithSpec :: forall a. HasSpec a => Specification a -> a -> Maybe a
+fixupWithSpec spec a
+  | a `conformsToSpec` spec = Just a
+  | otherwise = case spec of
+      MemberSpec (a' :| _) -> Just a'
+      TypeSpec ts _ -> fixupWithTypeSpec ts a >>= \ a' -> a' <$  guard (conformsToSpec a' spec)
+      _ -> listToMaybe $ filter (`conformsToSpec` spec) (shrinkWithSpec TrueSpec a)
+
 -- Debugging --------------------------------------------------------------
 
 -- | A version of `genFromSpecT` that runs in the IO monad. Good for debugging.
@@ -197,6 +254,10 @@ prettyPlan (simplifySpec -> spec)
 
 -- ---------------------- Building a plan -----------------------------------
 
+unsafeSubstStage :: Env -> SolverStage -> SolverStage
+unsafeSubstStage env (SolverStage y ps spec relevant) =
+  normalizeSolverStage $ SolverStage y (substPred env <$> ps) spec relevant
+
 substStage :: HasSpec a => Set Name -> Var a -> a -> SolverStage -> SolverStage
 substStage rel' x val (SolverStage y ps spec relevant) =
   normalizeSolverStage $ SolverStage y (substPred env <$> ps) spec relevant'
@@ -1105,6 +1166,9 @@ instance (HasSpec a, HasSpec b, KnownNat (CountCases b)) => HasSpec (Sum a b) wh
   shrinkWithTypeSpec (SumSpec _ sa _) (SumLeft a) = SumLeft <$> shrinkWithSpec sa a
   shrinkWithTypeSpec (SumSpec _ _ sb) (SumRight b) = SumRight <$> shrinkWithSpec sb b
 
+  fixupWithTypeSpec (SumSpec _ sa _) (SumLeft a) = SumLeft <$> fixupWithSpec sa a
+  fixupWithTypeSpec (SumSpec _ _ sb) (SumRight b) = SumRight <$> fixupWithSpec sb b
+
   toPreds ct (SumSpec h sa sb) =
     Case
       ct

src/Constrained/NumOrd.hs

@@ -39,6 +39,7 @@ module Constrained.NumOrd (
   combineNumSpec,
   genFromNumSpec,
   shrinkWithNumSpec,
+  fixupWithTypeSpec,
   conformsToNumSpec,
   toPredsNumSpec,
   OrdLike (..),
@@ -320,6 +321,11 @@ genFromNumSpec (NumSpecInterval ml mu) = do
 shrinkWithNumSpec :: Arbitrary n => NumSpec n -> n -> [n]
 shrinkWithNumSpec _ = shrink
 
+-- TODO: fixme
+
+fixupWithNumSpec :: Arbitrary n => NumSpec n -> n -> Maybe n
+fixupWithNumSpec _ = listToMaybe . shrink
+
 constrainInterval ::
   (MonadGenError m, Ord a, Num a, Show a) => Maybe a -> Maybe a -> Integer -> m (a, a)
 constrainInterval ml mu r =
@@ -1073,6 +1079,7 @@ instance HasSpec Integer where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec = cardinalNumSpec
@@ -1084,6 +1091,7 @@ instance HasSpec Int where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec = cardinalNumSpec
@@ -1095,6 +1103,7 @@ instance HasSpec (Ratio Integer) where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec _ = TrueSpec
@@ -1106,6 +1115,7 @@ instance HasSpec Natural where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec (NumSpecInterval (fromMaybe 0 -> lo) (Just hi)) =
@@ -1121,6 +1131,7 @@ instance HasSpec Word8 where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec = cardinalNumSpec
@@ -1134,6 +1145,7 @@ instance HasSpec Word16 where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec = cardinalNumSpec
@@ -1146,6 +1158,7 @@ instance HasSpec Word32 where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec = cardinalNumSpec
@@ -1157,6 +1170,7 @@ instance HasSpec Word64 where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec = cardinalNumSpec
@@ -1168,6 +1182,7 @@ instance HasSpec Int8 where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTrueSpec = equalSpec 256
@@ -1180,6 +1195,7 @@ instance HasSpec Int16 where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec = cardinalNumSpec
@@ -1192,6 +1208,7 @@ instance HasSpec Int32 where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec = cardinalNumSpec
@@ -1203,6 +1220,7 @@ instance HasSpec Int64 where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec = cardinalNumSpec
@@ -1214,6 +1232,7 @@ instance HasSpec Float where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec _ = TrueSpec
@@ -1225,6 +1244,7 @@ instance HasSpec Double where
   combineSpec = combineNumSpec
   genFromTypeSpec = genFromNumSpec
   shrinkWithTypeSpec = shrinkWithNumSpec
+  fixupWithTypeSpec = fixupWithNumSpec
   conformsTo = conformsToNumSpec
   toPreds = toPredsNumSpec
   cardinalTypeSpec _ = TrueSpec

src/Constrained/Spec/List.hs

@@ -244,6 +244,9 @@ instance HasSpec a => HasSpec [a] where
   shrinkWithTypeSpec (ListSpec _ _ _ es _) as =
     shrinkList (shrinkWithSpec es) as
 
+  -- TODO: fixme
+  fixupWithTypeSpec _ _ = Nothing
+
   cardinalTypeSpec _ = TrueSpec
 
   guardTypeSpec = guardListSpec

src/Constrained/Spec/Map.hs

@@ -278,6 +278,8 @@ instance
 
   shrinkWithTypeSpec (MapSpec _ _ _ _ kvs _) m = map Map.fromList $ shrinkList (shrinkWithSpec kvs) (Map.toList m)
 
+  fixupWithTypeSpec _ _ = Nothing
+
   toPreds m (MapSpec mHint mustKeys mustVals size kvs foldSpec) =
     toPred
       [ Assert $ Lit mustKeys `subset_` dom_ m

src/Constrained/Spec/Set.hs

@@ -204,6 +204,9 @@ instance (Ord a, HasSpec a) => HasSpec (Set a) where
 
   shrinkWithTypeSpec (SetSpec _ es _) as = map Set.fromList $ shrinkList (shrinkWithSpec es) (Set.toList as)
 
+  -- TODO: fixme
+  fixupWithTypeSpec _ _ = Nothing
+
   toPreds s (SetSpec m es size) =
     fold $
       -- Don't include this if the must set is empty

src/Constrained/Spec/Tree.hs

@@ -109,6 +109,9 @@ instance HasSpec a => HasSpec (Tree a) where
          | ts' <- shrinkList (shrinkWithTypeSpec (TreeSpec Nothing Nothing TrueSpec ctxSpec)) ts
          ]
 
+  -- TODO: fixme
+  fixupWithTypeSpec _ _ = Nothing
+
   cardinalTypeSpec _ = mempty
 
   toPreds t (TreeSpec mal msz rs s) =

src/Constrained/TheKnot.hs

@@ -136,6 +136,9 @@ instance (HasSpec a, HasSpec b) => HasSpec (Prod a b) where
     [Prod a' b | a' <- shrinkWithSpec sa a]
       ++ [Prod a b' | b' <- shrinkWithSpec sb b]
 
+  fixupWithTypeSpec (Cartesian sa sb) (Prod a b) =
+    Prod <$> fixupWithSpec sa a <*> fixupWithSpec sb b
+
   toPreds x (Cartesian sf ss) =
     satisfies (prodFst_ x) sf
       <> satisfies (prodSnd_ x) ss