input-output-hk / constrained-generators / 500

Committed 08 Jan 2026 10:49AM UTC coverage: 77.318% (+0.1%) from 77.177%

Build # 500

Build Type

push

github

Committed by

web-flow

Commit Message

Test and fix for bad behaviour of chooseSpec (#68)

Run Details

32 of 34 new or added lines in 4 files covered. (94.12%)

2 existing lines in 2 files now uncovered.

4019 of 5198 relevant lines covered (77.32%)

1.46 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

78.12

/src/Constrained/Generation.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- | All the things that are necessary for generation and shrinking.
module Constrained.Generation (
  -- * Generation and shrinking
  genFromSpec,
  genFromSpecT,
  genFromSpecWithSeed,
  shrinkWithSpec,
  fixupWithSpec,
  simplifySpec,

  -- ** Debugging
  printPlan,
  debugSpec,
  prettyPlan,

  -- * Function Symbols
  or_,
  not_,
  injRight_,
  injLeft_,
  (==.),

  -- * Other syntax
  whenTrue,

  -- * Internals
  CountCases,
  SumW (..),
  BoolW (..),
  EqW (..),
  SumSpec (..),
  pattern SumSpec,
  mapSpec,
  forwardPropagateSpec,
) where

import Constrained.AbstractSyntax
import Constrained.Base
import Constrained.Conformance
import Constrained.Core
import Constrained.Env (Env)
import Constrained.Env qualified as Env
import Constrained.FunctionSymbol
import Constrained.GenT
import Constrained.Generic
import Constrained.Graph hiding (irreflexiveDependencyOn)
import Constrained.List
import Constrained.NumOrd
import Constrained.PrettyUtils
import Constrained.Syntax
import Control.Applicative
import Control.Monad
import Control.Monad.Writer (Writer, runWriter, tell)
import Data.Foldable
import Data.Int
import Data.Kind
import Data.List (partition)
import Data.List.NonEmpty qualified as NE
import Data.Maybe
import Data.Semigroup (Any (..), getSum)
import Data.Semigroup qualified as Semigroup
import Data.Set (Set)
import Data.Set qualified as Set
import Data.String
import Data.Typeable
import GHC.Stack
import GHC.TypeLits
import Prettyprinter hiding (cat)
import Test.QuickCheck hiding (Args, Fun, Witness, forAll, witness)
import Test.QuickCheck.Gen
import Test.QuickCheck.Random hiding (left, right)
import Prelude hiding (cycle, pred)

------------------------------------------------------------------------
-- Generation, shrinking, and debugging
------------------------------------------------------------------------

-- | Generate a value that satisfies the spec. This function can fail if the
-- spec is inconsistent, there is a dependency error, or if the underlying
-- generators are not flexible enough.
genFromSpecT ::
  forall a m. (HasCallStack, HasSpec a, MonadGenError m) => Specification a -> GenT m a
genFromSpecT (ExplainSpec [] s) = genFromSpecT s
genFromSpecT (ExplainSpec es s) = push es (genFromSpecT s)
genFromSpecT (simplifySpec -> spec) = case spec of
  ExplainSpec [] s -> genFromSpecT s
  ExplainSpec es s -> push es (genFromSpecT s)
  MemberSpec as -> explain ("genFromSpecT on spec" ++ show spec) $ pureGen (elements (NE.toList as))
  TrueSpec -> genFromSpecT (typeSpec $ emptySpec @a)
  SuspendedSpec x p
    -- NOTE: If `x` isn't free in `p` we still have to try to generate things
    -- from `p` to make sure `p` is sat and then we can throw it away. A better
    -- approach would be to only do this in the case where we don't know if `p`
    -- is sat. The proper way to implement such a sat check is to remove
    -- sat-but-unnecessary variables in the optimiser.
    | not $ Name x `appearsIn` p -> do
        !_ <- genFromPreds mempty p
        genFromSpecT TrueSpec
    | otherwise -> do
        env <- genFromPreds mempty p
        Env.find env x
  TypeSpec s cant -> do
    mode <- getMode
    explainNE
      ( NE.fromList
          [ "genFromSpecT on (TypeSpec tspec cant) at type " ++ showType @a
          , "tspec = "
          , show s
          , "cant = " ++ show cant
          , "with mode " ++ show mode
          ]
      )
      $
      -- TODO: we could consider giving `cant` as an argument to `genFromTypeSpec` if this
      -- starts giving us trouble.
      genFromTypeSpec s `suchThatT` (`notElem` cant)
  ErrorSpec e -> genErrorNE e

-- | A version of `genFromSpecT` that simply errors if the generator fails
genFromSpec :: forall a. (HasCallStack, HasSpec a) => Specification a -> Gen a
genFromSpec spec = do
  res <- catchGen $ genFromSpecT @a @GE spec
  either (error . ('\n' :) . catMessages) pure res

-- | A version of `genFromSpecT` that takes a seed and a size and gives you a result
genFromSpecWithSeed ::
  forall a. (HasCallStack, HasSpec a) => Int -> Int -> Specification a -> a
genFromSpecWithSeed seed size spec = unGen (genFromSpec spec) (mkQCGen seed) size

-- ----------------------- Shrinking -------------------------------

unconstrainedShrink :: forall a. HasSpec a => a -> [a]
unconstrainedShrink = shrinkWithTypeSpec (emptySpec @a)

-- | Shrink a value while preserving adherence to a `Specification`
shrinkWithSpec :: forall a. HasSpec a => Specification a -> a -> [a]
shrinkWithSpec (ExplainSpec _ s) a = shrinkWithSpec s a
shrinkWithSpec (simplifySpec -> spec) a = case spec of
  -- TODO: It would be nice to avoid the extra `conformsToSpec` check here and only look
  -- at the cant set instead
  TypeSpec s _ -> [a' | a' <- shrinkWithTypeSpec s a, a' `conformsToSpec` spec]
  SuspendedSpec x p -> shrinkFromPreds p x a
  -- TODO: it would be nice if there was some better way of doing this
  MemberSpec as -> [a' | a' <- unconstrainedShrink a, a' `elem` as]
  TrueSpec -> unconstrainedShrink a
  ErrorSpec {} -> []
  -- Should be impossible?
  ExplainSpec _ s -> shrinkWithSpec s a

shrinkFromPreds :: forall a. 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!"
      if not $ Name x `appearsIn` p -- NOTE: this is safe because we just checked that p is SAT above
        then return $ unconstrainedShrink a
        else do
          -- 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.
debugSpec :: forall a. HasSpec a => Specification a -> IO ()
debugSpec spec = do
  ans <- generate $ genFromGenT $ inspect (genFromSpecT spec)
  let f x = putStrLn (unlines (NE.toList x))
      ok x =
        if conformsToSpec x spec
          then putStrLn "True"
          else putStrLn "False, perhaps there is an unsafeExists in the spec?"
  case ans of
    FatalError xs -> mapM_ f xs
    GenError xs -> mapM_ f xs
    Result x -> print spec >> print x >> ok x

-- | Pretty-print the plan for a `Specifcation` in the terminal for debugging
printPlan :: HasSpec a => Specification a -> IO ()
printPlan = print . prettyPlan

-- | Plan pretty-printer for debugging
prettyPlan :: HasSpec a => Specification a -> Doc ann
prettyPlan (simplifySpec -> spec)
  | SuspendedSpec _ p <- spec
  , Result plan <- prepareLinearization p =
      vsep'
        [ "Simplified spec:" /> pretty spec
        , pretty plan
        ]
  | otherwise = "Simplfied spec:" /> pretty 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'
  where
    env = Env.singleton x val
    relevant'
      | Name x `appearsIn` ps = rel' <> relevant
      | otherwise = relevant

normalizeSolverStage :: SolverStage -> SolverStage
normalizeSolverStage (SolverStage x ps spec relevant) = SolverStage x ps'' (spec <> spec') relevant
  where
    (ps', ps'') = partition ((1 ==) . Set.size . freeVarSet) ps
    spec' = fromGESpec $ computeSpec x (And ps')

-- TODO: here we can compute both the explicit hints (i.e. constraints that
-- define the order of two variables) and any whole-program smarts.
computeHints :: [Pred] -> Hints
computeHints ps =
  transitiveClosure $ fold [x `irreflexiveDependencyOn` y | DependsOn x y <- ps]

-- | Linearize a predicate, turning it into a list of variables to solve and
-- their defining constraints such that each variable can be solved independently.
prepareLinearization :: Pred -> GE SolverPlan
prepareLinearization p = do
  let preds = concatMap saturatePred $ flattenPred p
      hints = computeHints preds
      graph = transitiveClosure $ hints <> respecting hints (foldMap computeDependencies preds)
  plan <-
    explainNE
      ( NE.fromList
          [ "Linearizing"
          , show $
              "  preds: "
                <> pretty (take 3 preds)
                <> (if length preds > 3 then fromString (" ... " ++ show (length preds - 3) ++ " more.") else "")
          , show $ "  graph: " <> pretty graph
          ]
      )
      $ linearize preds graph
  pure $ backPropagation mempty $ SolverPlan plan

-- | Flatten nested `Let`, `Exists`, and `And` in a `Pred fn`. `Let` and
-- `Exists` bound variables become free in the result.
flattenPred :: Pred -> [Pred]
flattenPred pIn = go (freeVarNames pIn) [pIn]
  where
    go _ [] = []
    go fvs (p : ps) = case p of
      And ps' -> go fvs (ps' ++ ps)
      -- NOTE: the order of the arguments to `==.` here are important.
      -- The whole point of `Let` is that it allows us to solve all of `t`
      -- before we solve the variables in `t`.
      Let t b -> goBinder fvs b ps (\x -> (assert (t ==. (V x)) :))
      Exists _ b -> goBinder fvs b ps (const id)
      When b pp -> map (When b) (go fvs [pp]) ++ go fvs ps
      Explain es pp -> map (explanation es) (go fvs [pp]) ++ go fvs ps
      _ -> p : go fvs ps

    goBinder ::
      Set Int ->
      Binder a ->
      [Pred] ->
      (HasSpec a => Var a -> [Pred] -> [Pred]) ->
      [Pred]
    goBinder fvs (x :-> p) ps k = k x' $ go (Set.insert (nameOf x') fvs) (p' : ps)
      where
        (x', p') = freshen x p fvs

-- Consider: A + B = C + D
-- We want to fail if A and B are independent.
-- Consider: A + B = A + C, A <- B
-- Here we want to consider this constraint defining for A
linearize ::
  MonadGenError m => [Pred] -> DependGraph -> m [SolverStage]
linearize preds graph = do
  sorted <- case topsort graph of
    Left cycle ->
      fatalError
        ( show $
            "linearize: Dependency cycle in graph:"
              /> vsep'
                [ "cycle:" /> pretty cycle
                , "graph:" /> pretty graph
                ]
        )
    Right sorted -> pure sorted
  go sorted [(freeVarSet ps, ps) | ps <- filter isRelevantPred preds]
  where
    isRelevantPred TruePred = False
    isRelevantPred DependsOn {} = False
    isRelevantPred (Assert (Lit True)) = False
    isRelevantPred _ = True

    go [] [] = pure []
    go [] ps
      | null $ foldMap fst ps =
          case checkPredsE (pure "Linearizing fails") mempty (map snd ps) of
            Nothing -> pure []
            Just msgs -> genErrorNE msgs
      | otherwise =
          fatalErrorNE $
            NE.fromList
              [ "Dependency error in `linearize`: "
              , show $ indent 2 $ "graph: " /> pretty graph
              , show $
                  indent 2 $
                    "the following left-over constraints are not defining constraints for a unique variable:"
                      /> vsep' (map (pretty . snd) ps)
              ]
    go (n@(Name x) : ns) ps = do
      let (nps, ops) = partition (isLastVariable n . fst) ps
      (normalizeSolverStage (SolverStage x (map snd nps) mempty mempty) :) <$> go ns ops

    isLastVariable n set = n `Set.member` set && solvableFrom n (Set.delete n set) graph

------------------------------------------------------------------------
-- Simplification of Specifications
------------------------------------------------------------------------

-- | Spec simplification, use with care and don't modify the spec after using this!
simplifySpec :: HasSpec a => Specification a -> Specification a
simplifySpec spec = case applyNameHints spec of
  SuspendedSpec x p ->
    let optP = optimisePred p
     in fromGESpec $
          explain
            ("\nWhile calling simplifySpec on var " ++ show x ++ "\noptP=\n" ++ show optP ++ "\n")
            (computeSpecSimplified x optP)
  MemberSpec xs -> MemberSpec xs
  ErrorSpec es -> ErrorSpec es
  TypeSpec ts cant -> TypeSpec ts cant
  TrueSpec -> TrueSpec
  ExplainSpec es s -> explainSpec es (simplifySpec s)

-- ------- Stages of simplifying -------------------------------

-- TODO: it might be necessary to run aggressiveInlining again after the let floating etc.
optimisePred :: Pred -> Pred
optimisePred p =
  simplifyPred
    . letSubexpressionElimination
    . letFloating
    . aggressiveInlining
    . simplifyPred
    $ p

aggressiveInlining :: Pred -> Pred
aggressiveInlining pred
  | inlined = aggressiveInlining pInlined
  | otherwise = pred
  where
    (pInlined, Any inlined) = runWriter $ go (freeVars pred) [] pred

    underBinder fvs x p = fvs `without` [Name x] <> singleton (Name x) (countOf (Name x) p)

    underBinderSub :: HasSpec a => Subst -> Var a -> Subst
    underBinderSub sub x =
      [ x' := t
      | x' := t <- sub
      , isNothing $ eqVar x x'
      ]

    -- NOTE: this is safe because we only use the `Subst` when it results in a literal so there
    -- is no risk of variable capture.
    goBinder :: FreeVars -> Subst -> Binder a -> Writer Any (Binder a)
    goBinder fvs sub (x :-> p) = (x :->) <$> go (underBinder fvs x p) (underBinderSub sub x) p

    -- Check that the name `n` is only ever used as the only variable
    -- in the expressions where it appears. This ensures that it doesn't
    -- interact with anything.
    onlyUsedUniquely n p = case p of
      Assert t
        | n `appearsIn` t -> Set.size (freeVarSet t) == 1
        | otherwise -> True
      And ps -> all (onlyUsedUniquely n) ps
      -- TODO: we can (and should) probably add a bunch of cases to this.
      _ -> False

    go fvs sub pred2 = case pred2 of
      ElemPred bool t xs
        | not (isLit t)
        , Lit a <- substituteAndSimplifyTerm sub t -> do
            tell $ Any True
            pure $ ElemPred bool (Lit a) xs
        | otherwise -> pure $ ElemPred bool t xs
      Subst x t p -> go fvs sub (substitutePred x t p)
      Reifies t' t f
        | not (isLit t)
        , Lit a <- substituteAndSimplifyTerm sub t -> do
            tell $ Any True
            pure $ Reifies t' (Lit a) f
        | otherwise -> pure $ Reifies t' t f
      ForAll set b
        | not (isLit set)
        , Lit a <- substituteAndSimplifyTerm sub set -> do
            tell $ Any True
            pure $ foldMap (`unBind` b) (forAllToList a)
        | otherwise -> ForAll set <$> goBinder fvs sub b
      Case t bs
        | not (isLit t)
        , Lit a <- substituteAndSimplifyTerm sub t -> do
            tell $ Any True
            pure $ runCaseOn a (mapList thing bs) $ \x v p -> substPred (Env.singleton x v) p
        | (Weighted w (x :-> p) :> Nil) <- bs -> do
            let t' = substituteAndSimplifyTerm sub t
            p' <- go (underBinder fvs x p) (x := t' : sub) p
            pure $ Case t (Weighted w (x :-> p') :> Nil)
        | otherwise -> Case t <$> mapMList (traverseWeighted $ goBinder fvs sub) bs
      When b tp
        | not (isLit b)
        , Lit a <- substituteAndSimplifyTerm sub b -> do
            tell $ Any True
            pure $ if a then tp else TruePred
        | otherwise -> whenTrue b <$> go fvs sub tp
      Let t (x :-> p)
        | all (\n -> count n fvs <= 1) (freeVarSet t) -> do
            tell $ Any True
            pure $ substitutePred x t p
        | onlyUsedUniquely (Name x) p -> do
            tell $ Any True
            pure $ substitutePred x t p
        | not $ Name x `appearsIn` p -> do
            tell $ Any True
            pure p
        | not (isLit t)
        , Lit a <- substituteAndSimplifyTerm sub t -> do
            tell $ Any True
            pure $ unBind a (x :-> p)
        | otherwise -> Let t . (x :->) <$> go (underBinder fvs x p) (x := t : sub) p
      Exists k b -> Exists k <$> goBinder fvs sub b
      And ps -> fold <$> mapM (go fvs sub) ps
      Assert t
        | not (isLit t)
        , Lit b <- substituteAndSimplifyTerm sub t -> do
            tell $ Any True
            pure $ toPred b
        | otherwise -> pure pred2
      -- If the term turns into a literal, there is no more generation to do here
      -- so we can ignore the `GenHint`
      GenHint _ t
        | not (isLit t)
        , Lit {} <- substituteAndSimplifyTerm sub t -> do
            tell $ Any True
            pure TruePred
        | otherwise -> pure pred2
      DependsOn t t'
        | not (isLit t)
        , Lit {} <- substituteAndSimplifyTerm sub t -> do
            tell $ Any True
            pure $ TruePred
        | not (isLit t')
        , Lit {} <- substituteAndSimplifyTerm sub t' -> do
            tell $ Any True
            pure $ TruePred
        | otherwise -> pure pred2
      TruePred -> pure pred2
      FalsePred {} -> pure pred2
      Monitor {} -> pure pred2
      Explain es p -> Explain es <$> go fvs sub p

-- | Apply a substitution and simplify the resulting term if the
-- substitution changed the term.
substituteAndSimplifyTerm :: Subst -> Term a -> Term a
substituteAndSimplifyTerm sub t =
  case runWriter $ substituteTerm' sub t of
    (t', Any b)
      | b -> simplifyTerm t'
      | otherwise -> t'

-- | Simplify a Term, if the Term is an 'App', apply the rewrite rules
--   chosen by the (Logic sym t bs a) instance attached
--   to the function witness 'f'
simplifyTerm :: forall a. Term a -> Term a
simplifyTerm = \case
  V v -> V v
  Lit l -> Lit l
  App (f :: t bs a) (mapList simplifyTerm -> ts)
    | Just vs <- fromLits ts -> Lit $ uncurryList_ unValue (semantics f) vs
    | Just t <- rewriteRules f ts (Evidence @(AppRequires t bs a)) -> simplifyTerm t
    | otherwise -> App f ts

simplifyPred :: Pred -> Pred
simplifyPred = \case
  -- If the term simplifies away to a literal, that means there is no
  -- more generation to do so we can get rid of `GenHint`
  GenHint h t -> case simplifyTerm t of
    Lit {} -> TruePred
    t' -> GenHint h t'
  p@(ElemPred bool t xs) -> case simplifyTerm t of
    Lit x -> case (elem x xs, bool) of
      (True, True) -> TruePred
      (True, False) -> FalsePred ("notElemPred reduces to True" :| [show p])
      (False, True) -> FalsePred ("elemPred reduces to False" :| [show p])
      (False, False) -> TruePred
    t' -> ElemPred bool t' xs
  Subst x t p -> simplifyPred $ substitutePred x t p
  Assert t -> Assert $ simplifyTerm t
  Reifies t' t f -> case simplifyTerm t of
    Lit a ->
      -- Assert $ simplifyTerm t' ==. Lit (f a)
      ElemPred True (simplifyTerm t') (pure (f a))
    t'' -> Reifies (simplifyTerm t') t'' f
  ForAll (ts :: Term t) (b :: Binder a) -> case simplifyTerm ts of
    Lit as -> foldMap (`unBind` b) (forAllToList as)
    -- (App (extractW (UnionW @t) -> Just Refl) xs) -> error "MADE IT"
    {- Has to wait until we have HasSpec(Set a) instance
    UnionPat (xs :: Term (Set a)) ys ->
       let b' = simplifyBinder b
       in mkForAll xs b' <> mkForAll ys b' -}
    set' -> case simplifyBinder b of
      _ :-> TruePred -> TruePred
      b' -> ForAll set' b'
  DependsOn _ Lit {} -> TruePred
  DependsOn Lit {} _ -> TruePred
  DependsOn x y -> DependsOn x y
  -- Here is where we need the SumSpec instance
  Case t bs
    | Just es <- buildElemList bs -> ElemPred True (simplifyTerm t) es
    | otherwise -> mkCase (simplifyTerm t) (mapList (mapWeighted simplifyBinder) bs)
    where
      buildElemList :: List (Weighted Binder) as -> Maybe (NE.NonEmpty (SumOver as))
      buildElemList Nil = Nothing
      buildElemList (Weighted Nothing (x :-> ElemPred True (V x') as) :> xs)
        | Just Refl <- eqVar x x' =
            case xs of
              Nil -> Just as
              _ :> _ -> do
                rest <- buildElemList xs
                return $ fmap SumLeft as <> fmap SumRight rest
      buildElemList _ = Nothing
  When b p -> whenTrue (simplifyTerm b) (simplifyPred p)
  TruePred -> TruePred
  FalsePred es -> FalsePred es
  And ps -> fold (simplifyPreds ps)
  Let t b -> case simplifyTerm t of
    t'@App {} -> Let t' (simplifyBinder b)
    -- Variable or literal
    t' | x :-> p <- b -> simplifyPred $ substitutePred x t' p
  Exists k b -> case simplifyBinder b of
    _ :-> TruePred -> TruePred
    -- This is to get rid of exisentials like:
    -- `constrained $ \ x -> exists $ \ y -> [x ==. y, y + 2 <. 10]`
    x :-> p | Just t <- pinnedBy x p -> simplifyPred $ substitutePred x t p
    b' -> Exists k b'
  Monitor {} -> TruePred
  -- TODO: This is a bit questionable. On the one hand we could get rid of `Explain` here
  -- and just return `simplifyPred p` but doing so risks missing explanations when things
  -- do go wrong.
  Explain es p -> explanation es $ simplifyPred p

simplifyPreds :: [Pred] -> [Pred]
simplifyPreds = go [] . map simplifyPred
  where
    go acc [] = reverse acc
    go _ (FalsePred err : _) = [FalsePred err]
    go acc (TruePred : ps) = go acc ps
    go acc (p : ps) = go (p : acc) ps

simplifyBinder :: Binder a -> Binder a
simplifyBinder (x :-> p) = x :-> simplifyPred p

-- TODO: this can probably be cleaned up and generalized along with generalizing
-- to make sure we float lets in some missing cases.
letFloating :: Pred -> Pred
letFloating = fold . go []
  where
    goBlock ctx ps = goBlock' (freeVarNames ctx <> freeVarNames ps) ctx ps

    goBlock' :: Set Int -> [Pred] -> [Pred] -> [Pred]
    goBlock' _ ctx [] = ctx
    goBlock' fvs ctx (Let t (x :-> p) : ps) =
      -- We can do `goBlock'` here because we've already done let floating
      -- on the inner `p`
      [Let t (x' :-> fold (goBlock' (Set.insert (nameOf x') fvs) ctx (p' : ps)))]
      where
        (x', p') = freshen x p fvs
    goBlock' fvs ctx (And ps : ps') = goBlock' fvs ctx (ps ++ ps')
    goBlock' fvs ctx (p : ps) = goBlock' fvs (p : ctx) ps

    goExists ::
      HasSpec a =>
      [Pred] ->
      (Binder a -> Pred) ->
      Var a ->
      Pred ->
      [Pred]
    goExists ctx ex x (Let t (y :-> p))
      | not $ Name x `appearsIn` t =
          let (y', p') = freshen y p (Set.insert (nameOf x) $ freeVarNames p <> freeVarNames t)
           in go ctx (Let t (y' :-> ex (x :-> p')))
    goExists ctx ex x p = ex (x :-> p) : ctx

    pushExplain es (Let t (x :-> p)) = Let t (x :-> pushExplain es p)
    pushExplain es (And ps) = And (pushExplain es <$> ps)
    pushExplain es (Exists k (x :-> p)) =
      Exists (explainSemantics k) (x :-> pushExplain es p)
      where
        -- TODO: Unfortunately this is necessary on ghc 8.10.7
        explainSemantics ::
          forall a.
          ((forall b. Term b -> b) -> GE a) ->
          (forall b. Term b -> b) ->
          GE a
        explainSemantics k2 env = explainNE es $ k2 env
    -- TODO: possibly one wants to have a `Term` level explanation in case
    -- the `b` propagates to ErrorSpec for some reason?
    pushExplain es (When b p) = When b (pushExplain es p)
    pushExplain es p = explanation es p

    go ctx = \case
      ElemPred bool t xs -> ElemPred bool t xs : ctx
      And ps0 -> goBlock ctx (map letFloating ps0)
      Let t (x :-> p) -> goBlock ctx [Let t (x :-> letFloating p)]
      Exists k (x :-> p) -> goExists ctx (Exists k) x (letFloating p)
      Subst x t p -> go ctx (substitutePred x t p)
      Reifies t' t f -> Reifies t' t f : ctx
      Explain es p -> pushExplain es p : ctx
      -- TODO: float let through forall if possible
      ForAll t (x :-> p) -> ForAll t (x :-> letFloating p) : ctx
      -- TODO: float let through the cases if possible
      Case t bs -> Case t (mapList (mapWeighted (\(x :-> p) -> x :-> letFloating p)) bs) : ctx
      -- TODO: float let through if possible
      When b p -> When b (letFloating p) : ctx
      -- Boring cases
      Assert t -> Assert t : ctx
      GenHint h t -> GenHint h t : ctx
      DependsOn t t' -> DependsOn t t' : ctx
      TruePred -> TruePred : ctx
      FalsePred es -> FalsePred es : ctx
      Monitor m -> Monitor m : ctx

-- Common subexpression elimination but only on terms that are already let-bound.
letSubexpressionElimination :: Pred -> Pred
letSubexpressionElimination = go []
  where
    adjustSub :: HasSpec a => Var a -> Subst -> Subst
    adjustSub x sub =
      [ x' := t
      | x' := t <- sub
      , isNothing $ eqVar x x'
      , -- TODO: possibly freshen the binder where
      -- `x` appears instead?
      not $ Name x `appearsIn` t
      ]

    goBinder :: Subst -> Binder a -> Binder a
    goBinder sub (x :-> p) = x :-> go (adjustSub x sub) p

    go sub = \case
      ElemPred bool t xs -> ElemPred bool (backwardsSubstitution sub t) xs
      GenHint h t -> GenHint h (backwardsSubstitution sub t)
      And ps -> And (go sub <$> ps)
      Let t (x :-> p) -> Let t' (x :-> go (x := t' : sub') p)
        where
          t' = backwardsSubstitution sub t
          sub' = adjustSub x sub
      Exists k b -> Exists k (goBinder sub b)
      Subst x t p -> go sub (substitutePred x t p)
      Assert t -> Assert (backwardsSubstitution sub t)
      Reifies t' t f -> Reifies (backwardsSubstitution sub t') (backwardsSubstitution sub t) f
      -- NOTE: this is a tricky case. One possible thing to do here is to keep the old `DependsOn t t'`
      -- and have the new DependsOn if `backwardsSubstitution` changed something. With this semantics you
      -- risk running into unintuitive behaviour if you have something like:
      -- ```
      -- let x = y + z in
      --  {y + z `dependsOn` w
      --   assert $ w <. y + 2
      --   ...}
      -- ```
      -- This will be rewritten as:
      -- ```
      -- let x = y + z in
      --  {z `dependsOn` w
      --   assert $ w <. y + 2
      --   ...}
      -- ```
      -- which changes the dependency order of `w` and `y`. However, fixing
      -- this behaviour in turn makes it more difficult to detect when
      -- variables are no longer used after being substituted away - which
      -- blocks some other optimizations. As we strongly encourage users not to
      -- use `letBind` in their own code most users will never encounter this issue
      -- so the tradeoff is probably worth it.
      DependsOn t t' -> DependsOn (backwardsSubstitution sub t) (backwardsSubstitution sub t')
      ForAll t b -> ForAll (backwardsSubstitution sub t) (goBinder sub b)
      Case t bs -> Case (backwardsSubstitution sub t) (mapList (mapWeighted $ goBinder sub) bs)
      When b p -> When (backwardsSubstitution sub b) (go sub p)
      TruePred -> TruePred
      FalsePred es -> FalsePred es
      Monitor m -> Monitor m
      Explain es p -> Explain es $ go sub p

-- Turning Preds into Specifications. Here is where Propagation occurs ----

-- | Precondition: the `Pred` defines the `Var a`
-- Runs in `GE` in order for us to have detailed context on failure.
computeSpecSimplified ::
  forall a. (HasSpec a, HasCallStack) => Var a -> Pred -> GE (Specification a)
computeSpecSimplified x pred3 = localGESpec $ case simplifyPred pred3 of
  ElemPred True t xs -> propagateSpec (MemberSpec xs) <$> toCtx x t
  ElemPred False (t :: Term b) xs -> propagateSpec (TypeSpec @b (emptySpec @b) (NE.toList xs)) <$> toCtx x t
  Monitor {} -> pure mempty
  GenHint h t -> propagateSpec (giveHint h) <$> toCtx x t
  Subst x' t p' -> computeSpec x (substitutePred x' t p') -- NOTE: this is impossible as it should have gone away already
  TruePred -> pure mempty
  FalsePred es -> genErrorNE es
  And ps -> do
    spec <- fold <$> mapM (computeSpecSimplified x) ps
    case spec of
      ExplainSpec es (SuspendedSpec y ps') -> pure $ explainSpec es (SuspendedSpec y $ simplifyPred ps')
      SuspendedSpec y ps' -> pure $ SuspendedSpec y $ simplifyPred ps'
      s -> pure s
  Let t b -> pure $ SuspendedSpec x (Let t b)
  Exists k b -> pure $ SuspendedSpec x (Exists k b)
  Assert (Lit True) -> pure mempty
  Assert (Lit False) -> genError (show pred3)
  Assert t -> propagateSpec (equalSpec True) <$> toCtx x t
  ForAll (Lit s) b -> fold <$> mapM (\val -> computeSpec x $ unBind val b) (forAllToList s)
  ForAll t b -> do
    bSpec <- computeSpecBinderSimplified b
    propagateSpec (fromForAllSpec bSpec) <$> toCtx x t
  Case (Lit val) bs -> runCaseOn val (mapList thing bs) $ \va vaVal psa -> computeSpec x (substPred (Env.singleton va vaVal) psa)
  Case t branches -> do
    branchSpecs <- mapMList (traverseWeighted computeSpecBinderSimplified) branches
    propagateSpec (caseSpec (Just (showType @a)) branchSpecs) <$> toCtx x t
  When (Lit b) tp -> if b then computeSpecSimplified x tp else pure TrueSpec
  -- This shouldn't happen a lot of the time because when the body is trivial we mostly get rid of the `When` entirely
  When {} -> pure $ SuspendedSpec x pred3
  Reifies (Lit a) (Lit val) f
    | f val == a -> pure TrueSpec
    | otherwise ->
        pure $
          ErrorSpec (NE.fromList ["Value does not reify to literal: " ++ show val ++ " -/> " ++ show a])
  Reifies t' (Lit val) f ->
    propagateSpec (equalSpec (f val)) <$> toCtx x t'
  Reifies Lit {} _ _ ->
    fatalErrorNE $ NE.fromList ["Dependency error in computeSpec: Reifies", "  " ++ show pred3]
  Explain es p -> do
    -- In case things crash in here we want the explanation
    s <- pushGE (NE.toList es) (computeSpecSimplified x p)
    -- This is because while we do want to propagate `explanation`s into `SuspendedSpec`
    -- we probably don't want to propagate the full "currently simplifying xyz" explanation.
    case s of
      SuspendedSpec x2 p2 -> pure $ SuspendedSpec x2 (explanation es p2)
      _ -> pure $ addToErrorSpec es s
  -- Impossible cases that should be ruled out by the dependency analysis and linearizer
  DependsOn {} ->
    fatalErrorNE $
      NE.fromList
        [ "The impossible happened in computeSpec: DependsOn"
        , "  " ++ show x
        , show $ indent 2 (pretty pred3)
        ]
  Reifies {} ->
    fatalErrorNE $
      NE.fromList
        ["The impossible happened in computeSpec: Reifies", "  " ++ show x, show $ indent 2 (pretty pred3)]
  where
    -- We want `genError` to turn into `ErrorSpec` and we want `FatalError` to turn into `FatalError`
    localGESpec ge = case ge of
      (GenError xs) -> Result $ ErrorSpec (catMessageList xs)
      (FatalError es) -> FatalError es
      (Result v) -> Result v

-- | Precondition: the `Pred fn` defines the `Var a`.
--   Runs in `GE` in order for us to have detailed context on failure.
computeSpec ::
  forall a. (HasSpec a, HasCallStack) => Var a -> Pred -> GE (Specification a)
computeSpec x p = computeSpecSimplified x (simplifyPred p)

computeSpecBinderSimplified :: Binder a -> GE (Specification a)
computeSpecBinderSimplified (x :-> p) = computeSpecSimplified x p

-- | Turn a list of branches into a SumSpec. If all the branches fail return an ErrorSpec.
--   Note the requirement of HasSpec(SumOver).
caseSpec ::
  forall as.
  HasSpec (SumOver as) =>
  Maybe String ->
  List (Weighted (Specification)) as ->
  Specification (SumOver as)
caseSpec tString ss
  | allBranchesFail ss =
      ErrorSpec
        ( NE.fromList
            [ "When simplifying SumSpec, all branches in a caseOn" ++ sumType tString ++ " simplify to False."
            , show spec
            ]
        )
  | True = spec
  where
    spec = loop tString ss

    allBranchesFail :: forall as2. List (Weighted Specification) as2 -> Bool
    allBranchesFail Nil = error "The impossible happened in allBranchesFail"
    allBranchesFail (Weighted _ s :> Nil) = isErrorLike s
    allBranchesFail (Weighted _ s :> ss2@(_ :> _)) = isErrorLike s && allBranchesFail ss2

    loop ::
      forall as3.
      HasSpec (SumOver as3) =>
      Maybe String ->
      List (Weighted Specification) as3 ->
      Specification (SumOver as3)
    loop _ Nil = error "The impossible happened in caseSpec"
    loop _ (s :> Nil) = thing s
    loop mTypeString (s :> ss1@(_ :> _))
      | Evidence <- prerequisites @(SumOver as3) =
          (typeSpec $ SumSpecRaw mTypeString theWeights (thing s) (loop Nothing ss1))
      where
        theWeights =
          case (weight s, totalWeight ss1) of
            (Nothing, Nothing) -> Nothing
            (a, b) -> Just (fromMaybe 1 a, fromMaybe (lengthList ss1) b)

------------------------------------------------------------------------
-- SumSpec et al
------------------------------------------------------------------------

-- | The Specification for Sums.
data SumSpec a b
  = SumSpecRaw
      (Maybe String) -- A String which is the type of arg in (caseOn arg branch1 .. branchN)
      (Maybe (Int, Int))
      (Specification a)
      (Specification b)

-- | The "normal" view of t`SumSpec` that doesn't take weights into account
pattern SumSpec ::
  (Maybe (Int, Int)) -> (Specification a) -> (Specification b) -> SumSpec a b
pattern SumSpec a b c <- SumSpecRaw _ a b c
  where
    SumSpec a b c = SumSpecRaw Nothing a b c

{-# COMPLETE SumSpec #-}

sumType :: Maybe String -> String
sumType Nothing = ""
sumType (Just x) = " type=" ++ x

totalWeight :: List (Weighted f) as -> Maybe Int
totalWeight = fmap getSum . foldMapList (fmap Semigroup.Sum . weight)

-- =================================
-- Operations on Stages and Plans

-- | Does nothing if the variable is not in the plan already.
mergeSolverStage :: SolverStage -> [SolverStage] -> [SolverStage]
mergeSolverStage (SolverStage x ps spec relevant) plan =
  [ case eqVar x y of
      Just Refl ->
        normalizeSolverStage $
          SolverStage
            y
            (ps ++ ps')
            (spec <> spec')
            (relevant <> relevant')
      Nothing -> stage
  | stage@(SolverStage y ps' spec' relevant') <- plan
  ]

isEmptyPlan :: SolverPlan -> Bool
isEmptyPlan (SolverPlan plan) = null plan

stepPlan :: MonadGenError m => SolverPlan -> Env -> SolverPlan -> GenT m (Env, SolverPlan)
stepPlan _ env plan@(SolverPlan []) = pure (env, plan)
stepPlan (SolverPlan origStages) env (SolverPlan (stage@(SolverStage (x :: Var a) ps spec relevant) : pl)) = do
  let errorMessage =
        "Failed to step the plan"
          /> vsep
            [ "Relevant parts of the original plan:" //> pretty narrowedOrigPlan
            , "Already generated variables:" //> pretty narrowedEnv
            , "Current stage:" //> pretty stage
            ]
      relevant' = Set.insert (Name x) relevant
      narrowedOrigPlan = SolverPlan $ [st | st@(SolverStage v _ _ _) <- origStages, Name v `Set.member` relevant']
      narrowedEnv = Env.filterKeys env (\v -> nameOf v `Set.member` (Set.map (\(Name n) -> nameOf n) relevant'))
  explain (show errorMessage) $ do
    when (isErrorLike spec) $
      genError "The specification in the current stage is unsatisfiable, giving up."
    when (not $ null ps) $
      fatalError
        "Something went wrong and not all predicates have been discharged. Report this as a bug in Constrained.Generation"
    val <- genFromSpecT spec
    let env1 = Env.extend x val env
    pure (env1, backPropagation relevant' $ SolverPlan (substStage relevant' x val <$> pl))

-- | Generate a satisfying `Env` for a `p : Pred fn`. The `Env` contains values for
-- all the free variables in `flattenPred p`.
genFromPreds :: forall m. MonadGenError m => Env -> Pred -> GenT m Env
-- TODO: remove this once optimisePred does a proper fixpoint computation
genFromPreds env0 (optimisePred . optimisePred -> preds) = do
  -- NOTE: this is just lazy enough that the work of flattening,
  -- computing dependencies, and linearizing is memoized in
  -- properties that use `genFromPreds`.
  origPlan <- runGE $ prepareLinearization preds
  let go :: Env -> SolverPlan -> GenT m Env
      go env plan | isEmptyPlan plan = pure env
      go env plan = do
        (env', plan') <- stepPlan origPlan env plan
        go env' plan'
  go env0 origPlan

-- | Push as much information we can backwards through the plan.
backPropagation :: Set Name -> SolverPlan -> SolverPlan
backPropagation relevant (SolverPlan initplan) = SolverPlan (go [] (reverse initplan))
  where
    go :: [SolverStage] -> [SolverStage] -> [SolverStage]
    go acc [] = acc
    go acc (s@(SolverStage (x :: Var a) ps spec _) : plan) = go (s : acc) plan'
      where
        newStages = concatMap newStage ps
        plan' = foldr mergeSolverStage plan newStages

        -- Note use of the Term Pattern Equal
        newStage (Assert (Equal tl tr))
          | [Name xl] <- Set.toList $ freeVarSet tl
          , [Name xr] <- Set.toList $ freeVarSet tr
          , 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)
                  ]
                _ -> []
        newStage (Case e bs)
          | [Name xe] <- Set.toList $ freeVarSet e
          , Nothing <- eqVar x xe =
              [ SolverStage
                  xe
                  [Case e $ mapList mkBranch bs]
                  TrueSpec
                  (Set.insert (Name x) relevant) -- TODO: this is only true in the
                  -- case where we actually rule some
                  -- stuff out
              ]
          where
            mkBranch :: Weighted Binder x -> Weighted Binder x
            mkBranch (Weighted _ (xb :-> p))
              | Result spec' <- computeSpec x p
              , isErrorLike (spec <> spec') =
                  Weighted Nothing $ xb :-> toPred False
              | otherwise = Weighted Nothing (xb :-> TruePred)
        newStage _ = []

-- | Function symbols for `(==.)`
data EqW :: [Type] -> Type -> Type where
  EqualW :: (Eq a, HasSpec a) => EqW '[a, a] Bool

deriving instance Eq (EqW dom rng)

instance Show (EqW d r) where
  show EqualW = "==."

instance Syntax EqW where
  isInfix EqualW = True

instance Semantics EqW where
  semantics EqualW = (==)

instance Logic EqW where
  propagate f ctxt (ExplainSpec es s) = explainSpec es $ propagate f ctxt s
  propagate _ _ TrueSpec = TrueSpec
  propagate _ _ (ErrorSpec msgs) = ErrorSpec msgs
  propagate EqualW (HOLE :? Value x :> Nil) (SuspendedSpec v ps) =
    constrained $ \v' -> Let (App EqualW (v' :> Lit x :> Nil)) (v :-> ps)
  propagate EqualW (Value x :! Unary HOLE) (SuspendedSpec v ps) =
    constrained $ \v' -> Let (App EqualW (Lit x :> v' :> Nil)) (v :-> ps)
  propagate EqualW (HOLE :? Value s :> Nil) spec =
    caseBoolSpec spec $ \case
      True -> equalSpec s
      False -> notEqualSpec s
  propagate EqualW (Value s :! Unary HOLE) spec =
    caseBoolSpec spec $ \case
      True -> equalSpec s
      False -> notEqualSpec s

  rewriteRules EqualW (t :> t' :> Nil) Evidence
    | t == t' = Just $ lit True
    | otherwise = Nothing

  saturate EqualW (FromGeneric (InjLeft _) :> t :> Nil) = [toPreds t (SumSpec Nothing TrueSpec (ErrorSpec (pure "saturatePred")))]
  saturate EqualW (FromGeneric (InjRight _) :> t :> Nil) = [toPreds t (SumSpec Nothing (ErrorSpec (pure "saturatePred")) TrueSpec)]
  saturate _ _ = []

infix 4 ==.

-- | Equality on the constraint-level
(==.) :: HasSpec a => Term a -> Term a -> Term Bool
(==.) = appTerm EqualW

-- | Pattern version of `(==.)` for rewrite rules
pattern Equal ::
  forall b.
  () =>
  forall a.
  (b ~ Bool, Eq a, HasSpec a) =>
  Term a ->
  Term a ->
  Term b
pattern Equal x y <-
  ( App
      (getWitness -> Just EqualW)
      (x :> y :> Nil)
    )

-- | Like @if b then p else assert True@ in constraint-land
whenTrue :: forall p. IsPred p => Term Bool -> p -> Pred
whenTrue (Lit True) (toPred -> p) = p
whenTrue (Lit False) _ = TruePred
whenTrue b (toPred -> FalsePred {}) = assert (not_ b)
whenTrue _ (toPred -> TruePred) = TruePred
whenTrue b (toPred -> p) = When b p

-- | Is the variable x pinned to some free term in p? (free term
-- meaning that all the variables in the term are free in p).
--
-- TODO: complete this with more cases!
pinnedBy :: forall a. HasSpec a => Var a -> Pred -> Maybe (Term a)
pinnedBy x (Assert (Equal t t'))
  | V x' <- t, Just Refl <- eqVar x x' = Just t'
  | V x' <- t', Just Refl <- eqVar x x' = Just t
pinnedBy x (ElemPred True (V x') (xs NE.:| []))
  | Just Refl <- eqVar x x' = Just (lit xs)
pinnedBy x (And ps) = listToMaybe $ catMaybes $ map (pinnedBy x) ps
pinnedBy x (Explain _ p) = pinnedBy x p
pinnedBy _ _ = Nothing

-- ==================================================================================================
-- TODO: generalize this to make it more flexible and extensible
--
-- The idea here is that we turn constraints into _extra_ constraints. C.f. the
-- `mapIsJust` example in `Constrained.Examples.Map`:

--    mapIsJust :: Specification BaseFn (Int, Int)
--    mapIsJust = constrained' $ \ [var| x |] [var| y |] ->
--      assert $ just_ x ==. lookup_ y (lit $ Map.fromList [(z, z) | z <- [100 .. 102]])

-- Without this code the example wouldn't work because `y` is completely unconstrained during
-- generation. With this code we _essentially_ rewrite occurences of `just_ A == B` to
-- `[just_ A == B, case B of Nothing -> False; Just _ -> True]` to add extra information
-- about the variables in `B`. Consequently, `y` in the example above is
-- constrained to `MemberSpec [100 .. 102]` in the plan. This is implemented using the `saturate`
-- function in the logic type class - in the example above for `==`.
saturatePred :: Pred -> [Pred]
saturatePred p =
  -- [p]
  --  + ---- if there is an Explain, it is still on 'p' here
  --  |
  --  v
  p : case p of
    Explain _es x -> saturatePred x -- Note that the Explain is still on the original 'p', so it is not lost
    -- Note how the saturation is done by the 'saturate' method of the Logic class
    Assert (App sym xs) -> saturate sym xs
    _ -> []

-- ==================================================================
-- HasSpec for Sums
-- ==================================================================

guardSumSpec ::
  forall a b.
  (HasSpec a, HasSpec b, KnownNat (CountCases b)) =>
  [String] ->
  SumSpec a b ->
  Specification (Sum a b)
guardSumSpec msgs s@(SumSpecRaw tString _ sa sb)
  | isErrorLike sa
  , isErrorLike sb =
      ErrorSpec $
        NE.fromList $
          msgs ++ ["All branches in a caseOn" ++ sumType tString ++ " simplify to False.", show s]
  | otherwise = typeSpec s

instance (KnownNat (CountCases b), HasSpec a, HasSpec b) => Show (SumSpec a b) where
  show sumspec@(SumSpecRaw tstring hint l r) = case alternateShow @(Sum a b) sumspec of
    (BinaryShow _ ps) -> show $ parens (fromString ("SumSpec" ++ sumType tstring) /> vsep ps)
    NonBinary ->
      "(SumSpec"
        ++ sumType tstring
        ++ show (sumWeightL hint)
        ++ " ("
        ++ show l
        ++ ") "
        ++ show (sumWeightR hint)
        ++ " ("
        ++ show r
        ++ "))"

combTypeName :: Maybe String -> Maybe String -> Maybe String
combTypeName (Just x) (Just y) =
  if x == y then Just x else Just ("(" ++ x ++ " | " ++ y ++ ")")
combTypeName (Just x) Nothing = Just x
combTypeName Nothing (Just x) = Just x
combTypeName Nothing Nothing = Nothing

instance (HasSpec a, HasSpec b) => Semigroup (SumSpec a b) where
  SumSpecRaw t h sa sb <> SumSpecRaw t' h' sa' sb' =
    SumSpecRaw (combTypeName t t') (unionWithMaybe mergeH h h') (sa <> sa') (sb <> sb')
    where
      -- TODO: think more carefully about this, now weights like 2 2 and 10 15 give more weight to 10 15
      -- than would be the case if you had 2 2 and 2 3. But on the other hand this approach is associative
      -- whereas actually averaging the ratios is not. One could keep a list. Future work.
      mergeH (fA, fB) (fA', fB') = (fA + fA', fB + fB')

instance forall a b. (HasSpec a, HasSpec b, KnownNat (CountCases b)) => Monoid (SumSpec a b) where
  mempty = SumSpec Nothing mempty mempty

-- | How many constructors are there in this type?
type family CountCases a where
  CountCases (Sum a b) = 1 + CountCases b
  CountCases _ = 1

countCases :: forall a. KnownNat (CountCases a) => Int
countCases = fromIntegral (natVal @(CountCases a) Proxy)

-- | The HasSpec Sum instance
instance (HasSpec a, HasSpec b, KnownNat (CountCases b)) => HasSpec (Sum a b) where
  type TypeSpec (Sum a b) = SumSpec a b

  type Prerequisites (Sum a b) = (HasSpec a, HasSpec b)

  emptySpec = mempty

  combineSpec s s' = guardSumSpec ["When combining SumSpecs", "  " ++ show s, "  " ++ show s'] (s <> s')

  conformsTo (SumLeft a) (SumSpec _ sa _) = conformsToSpec a sa
  conformsTo (SumRight b) (SumSpec _ _ sb) = conformsToSpec b sb

  genFromTypeSpec (SumSpec h sa sb)
    | emptyA, emptyB = genError "genFromTypeSpec @SumSpec: empty"
    | emptyA = SumRight <$> genFromSpecT sb
    | emptyB = SumLeft <$> genFromSpecT sa
    | fA == 0, fB == 0 = genError "All frequencies 0"
    | otherwise =
        frequencyT
          [ (fA, SumLeft <$> genFromSpecT sa)
          , (fB, SumRight <$> genFromSpecT sb)
          ]
    where
      (max 0 -> fA, max 0 -> fB) = fromMaybe (1, countCases @b) h
      emptyA = isErrorLike sa
      emptyB = isErrorLike sb

  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
      ( (Weighted (fst <$> h) $ bind $ \a -> satisfies a sa)
          :> (Weighted (snd <$> h) $ bind $ \b -> satisfies b sb)
          :> Nil
      )

  cardinalTypeSpec (SumSpec _ leftspec rightspec) = addSpecInt (cardinality leftspec) (cardinality rightspec)

  typeSpecHasError (SumSpec _ x y) =
    case (isErrorLike x, isErrorLike y) of
      (True, True) -> Just $ (errorLikeMessage x <> errorLikeMessage y)
      _ -> Nothing

  alternateShow (SumSpec h left right@(TypeSpec r [])) =
    case alternateShow @b r of
      (BinaryShow "SumSpec" ps) -> BinaryShow "SumSpec" ("|" <+> sumWeightL h <+> viaShow left : ps)
      (BinaryShow "Cartesian" ps) ->
        BinaryShow "SumSpec" ("|" <+> sumWeightL h <+> viaShow left : [parens ("Cartesian" /> vsep ps)])
      _ ->
        BinaryShow "SumSpec" ["|" <+> sumWeightL h <+> viaShow left, "|" <+> sumWeightR h <+> viaShow right]
  alternateShow (SumSpec h left right) =
    BinaryShow "SumSpec" ["|" <+> sumWeightL h <+> viaShow left, "|" <+> sumWeightR h <+> viaShow right]

-- ======================================
-- Here are the Logic Instances for Sum

-- | Function symbols for `injLeft_` and `injRight_`
data SumW dom rng where
  InjLeftW :: SumW '[a] (Sum a b)
  InjRightW :: SumW '[b] (Sum a b)

instance Show (SumW dom rng) where
  show InjLeftW = "injLeft_"
  show InjRightW = "injRight_"

deriving instance (Eq (SumW dom rng))

instance Syntax SumW

instance Semantics SumW where
  semantics InjLeftW = SumLeft
  semantics InjRightW = SumRight

instance Logic SumW where
  propagateTypeSpec InjLeftW (Unary HOLE) (SumSpec _ sl _) cant = sl <> foldMap notEqualSpec [a | SumLeft a <- cant]
  propagateTypeSpec InjRightW (Unary HOLE) (SumSpec _ _ sr) cant = sr <> foldMap notEqualSpec [a | SumRight a <- cant]

  propagateMemberSpec InjLeftW (Unary HOLE) es =
    case [a | SumLeft a <- NE.toList es] of
      (x : xs) -> MemberSpec (x :| xs)
      [] ->
        ErrorSpec $
          pure $
            "propMemberSpec (sumleft_ HOLE) on (MemberSpec es) with no SumLeft in es: " ++ show (NE.toList es)
  propagateMemberSpec InjRightW (Unary HOLE) es =
    case [a | SumRight a <- NE.toList es] of
      (x : xs) -> MemberSpec (x :| xs)
      [] ->
        ErrorSpec $
          pure $
            "propagate(InjRight HOLE) on (MemberSpec es) with no SumLeft in es: " ++ show (NE.toList es)

  mapTypeSpec InjLeftW ts = typeSpec $ SumSpec Nothing (typeSpec ts) (ErrorSpec (pure "mapTypeSpec InjLeftW"))
  mapTypeSpec InjRightW ts = typeSpec $ SumSpec Nothing (ErrorSpec (pure "mapTypeSpec InjRightW")) (typeSpec ts)

-- | Constructor for `Sum`
injLeft_ :: (HasSpec a, HasSpec b, KnownNat (CountCases b)) => Term a -> Term (Sum a b)
injLeft_ = appTerm InjLeftW

-- | Constructor for `Sum`
injRight_ :: (HasSpec a, HasSpec b, KnownNat (CountCases b)) => Term b -> Term (Sum a b)
injRight_ = appTerm InjRightW

-- | Pattern for building custom rewrite rules
pattern InjRight ::
  forall c.
  () =>
  forall a b.
  ( c ~ Sum a b
  , AppRequires SumW '[b] c
  ) =>
  Term b ->
  Term c
pattern InjRight x <- (App (getWitness -> Just InjRightW) (x :> Nil))

-- | Pattern for building custom rewrite rules
pattern InjLeft ::
  forall c.
  () =>
  forall a b.
  ( c ~ Sum a b
  , AppRequires SumW '[a] c
  ) =>
  Term a ->
  Term c
pattern InjLeft x <- App (getWitness -> Just InjLeftW) (x :> Nil)

sumWeightL, sumWeightR :: Maybe (Int, Int) -> Doc a
sumWeightL Nothing = "1"
sumWeightL (Just (x, _)) = fromString (show x)
sumWeightR Nothing = "1"
sumWeightR (Just (_, x)) = fromString (show x)

-- | Operations on Bool
data BoolW (dom :: [Type]) (rng :: Type) where
  NotW :: BoolW '[Bool] Bool
  OrW :: BoolW '[Bool, Bool] Bool

deriving instance Eq (BoolW dom rng)

instance Show (BoolW dom rng) where
  show NotW = "not_"
  show OrW = "or_"

boolSem :: BoolW dom rng -> FunTy dom rng
boolSem NotW = not
boolSem OrW = (||)

instance Semantics BoolW where
  semantics = boolSem

instance Syntax BoolW

-- ======= Logic instance BoolW

instance Logic BoolW where
  propagate f ctxt (ExplainSpec [] s) = propagate f ctxt s
  propagate f ctxt (ExplainSpec es s) = ExplainSpec es $ propagate f ctxt s
  propagate _ _ TrueSpec = TrueSpec
  propagate _ _ (ErrorSpec msgs) = ErrorSpec msgs
  propagate NotW (Unary HOLE) (SuspendedSpec v ps) =
    constrained $ \v' -> Let (App NotW (v' :> Nil)) (v :-> ps)
  propagate NotW (Unary HOLE) spec =
    caseBoolSpec spec (equalSpec . not)
  propagate OrW (HOLE :<: x) (SuspendedSpec v ps) =
    constrained $ \v' -> Let (App OrW (v' :> Lit x :> Nil)) (v :-> ps)
  propagate OrW (x :>: HOLE) (SuspendedSpec v ps) =
    constrained $ \v' -> Let (App OrW (Lit x :> v' :> Nil)) (v :-> ps)
  propagate OrW (HOLE :<: s) spec =
    caseBoolSpec spec (okOr s)
  propagate OrW (s :>: HOLE) spec =
    caseBoolSpec spec (okOr s)

  mapTypeSpec NotW () = typeSpec ()

-- | We have something like ('constant' ||. HOLE) must evaluate to 'need'.
--   Return a (Specification Bool) for HOLE, that makes that True.
okOr :: Bool -> Bool -> Specification Bool
okOr constant need = case (constant, need) of
  (True, True) -> TrueSpec
  (True, False) ->
    ErrorSpec
      (pure ("(" ++ show constant ++ "||. HOLE) must equal False. That cannot be the case."))
  (False, False) -> MemberSpec (pure False)
  (False, True) -> MemberSpec (pure True)

-- | Disjunction on @`Term` `Bool`@, note that this will not cause backtracking during generation
or_ :: Term Bool -> Term Bool -> Term Bool
or_ = appTerm OrW

-- | Negation of booleans
not_ :: Term Bool -> Term Bool
not_ = appTerm NotW

-- ===============================================================================
-- Syntax for Solving : stages and plans

data SolverStage where
  SolverStage ::
    HasSpec a =>
    { stageVar :: Var a
    , stagePreds :: [Pred]
    , stageSpec :: Specification a
    , relevantVariables :: Set Name
    } ->
    SolverStage

docVar :: Typeable a => Var a -> Doc h
docVar (v :: Var a) = fromString (show v ++ " :: " ++ showType @a)

instance Pretty SolverStage where
  pretty SolverStage {..} =
    docVar stageVar
      <+> "<-"
        /> vsep'
          ( [pretty stageSpec | not $ isTrueSpec stageSpec]
              ++ ["---" | not $ null stagePreds, not $ isTrueSpec stageSpec]
              ++ map pretty stagePreds
              ++ ["_" | null stagePreds && isTrueSpec stageSpec]
          )

newtype SolverPlan = SolverPlan {solverPlan :: [SolverStage]}

instance Pretty SolverPlan where
  pretty SolverPlan {..} =
    "SolverPlan" /> prettyLinear solverPlan

isTrueSpec :: Specification a -> Bool
isTrueSpec TrueSpec = True
isTrueSpec _ = False

prettyLinear :: [SolverStage] -> Doc ann
prettyLinear = vsep' . map pretty

fromGESpec :: HasCallStack => GE (Specification a) -> Specification a
fromGESpec ge = case ge of
  Result s -> s
  GenError xs -> ErrorSpec (catMessageList xs)
  FatalError es -> error $ catMessages es

-- | 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

input-output-hk / constrained-generators / 500

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous