-- editorconfig-checker-disable-file

module PlutusCore.Generators.QuickCheck.Substitutions where

import PlutusCore.Generators.QuickCheck.Common
import PlutusCore.Generators.QuickCheck.GenerateTypes
import PlutusCore.Generators.QuickCheck.ShrinkTypes
import PlutusCore.Generators.QuickCheck.Utils

import PlutusCore.Default
import PlutusCore.Name.Unique
import PlutusIR
import PlutusIR.Subst

import Control.Monad ((<=<))
import Data.Map.Strict (Map)
import Data.Map.Strict.Internal qualified as Map
import Data.Maybe
import Data.MultiSet (toMap)
import Data.MultiSet.Lens
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Set.Lens
import GHC.Stack
import Test.QuickCheck hiding (reason)

type TypeSub = Map TyName (Type TyName DefaultUni ())

{- Note [Substitution in generators]
Generators define their own substitution algorithm. It's generally useful, however it's quite
inefficient, since calling 'renameVar' in the worker when handling bindings makes the whole
algorithm at least quadratic. Plus, it also very general, since it handles both nested and parallel
substitution.

So for now we've decided to keep it local to the generators and extract it into the main library
once there's a specific reason to do that.
-}

-- | The most general substitution worker.
substTypeCustomGo
    :: HasCallStack
    => Bool                       -- ^ Nested ('True') or parallel ('False')
    -> Set TyName                 -- ^ Variables that are considered free.
    -> TypeSub                    -- ^ Type substitution to use.
    -> Type TyName DefaultUni ()  -- ^ Type to substitute in.
    -> Type TyName DefaultUni ()
substTypeCustomGo :: HasCallStack =>
Bool
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
substTypeCustomGo Bool
nested Set TyName
fvs0 = Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go Set TyName
fvs0 Set TyName
forall a. Set a
Set.empty where
    go :: Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go Set TyName
fvs Set TyName
seen TypeSub
sub Type TyName DefaultUni ()
ty = case Type TyName DefaultUni ()
ty of
        TyVar ()
_ TyName
x | TyName -> Set TyName -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member TyName
x Set TyName
seen -> [Char] -> Type TyName DefaultUni ()
forall a. HasCallStack => [Char] -> a
error [Char]
"substType' loop"
        -- In the case where we do nested substitution we just continue, in parallel substitution
        -- we never go below a substitution.
        TyVar ()
_ TyName
x | Bool
nested    -> Type TyName DefaultUni ()
-> (Type TyName DefaultUni () -> Type TyName DefaultUni ())
-> Maybe (Type TyName DefaultUni ())
-> Type TyName DefaultUni ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Type TyName DefaultUni ()
ty (Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go Set TyName
fvs (TyName -> Set TyName -> Set TyName
forall a. Ord a => a -> Set a -> Set a
Set.insert TyName
x Set TyName
seen) TypeSub
sub) (Maybe (Type TyName DefaultUni ()) -> Type TyName DefaultUni ())
-> Maybe (Type TyName DefaultUni ()) -> Type TyName DefaultUni ()
forall a b. (a -> b) -> a -> b
$ TyName -> TypeSub -> Maybe (Type TyName DefaultUni ())
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TyName
x TypeSub
sub
                  | Bool
otherwise -> Type TyName DefaultUni ()
-> (Type TyName DefaultUni () -> Type TyName DefaultUni ())
-> Maybe (Type TyName DefaultUni ())
-> Type TyName DefaultUni ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Type TyName DefaultUni ()
ty Type TyName DefaultUni () -> Type TyName DefaultUni ()
forall a. a -> a
id (Maybe (Type TyName DefaultUni ()) -> Type TyName DefaultUni ())
-> Maybe (Type TyName DefaultUni ()) -> Type TyName DefaultUni ()
forall a b. (a -> b) -> a -> b
$ TyName -> TypeSub -> Maybe (Type TyName DefaultUni ())
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TyName
x TypeSub
sub
        TyFun ()
_ Type TyName DefaultUni ()
a Type TyName DefaultUni ()
b -> ()
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
forall tyname (uni :: * -> *) ann.
ann
-> Type tyname uni ann
-> Type tyname uni ann
-> Type tyname uni ann
TyFun () (Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go Set TyName
fvs Set TyName
seen TypeSub
sub Type TyName DefaultUni ()
a) (Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go Set TyName
fvs Set TyName
seen TypeSub
sub Type TyName DefaultUni ()
b)
        TyApp ()
_ Type TyName DefaultUni ()
a Type TyName DefaultUni ()
b -> ()
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
forall tyname (uni :: * -> *) ann.
ann
-> Type tyname uni ann
-> Type tyname uni ann
-> Type tyname uni ann
TyApp () (Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go Set TyName
fvs Set TyName
seen TypeSub
sub Type TyName DefaultUni ()
a) (Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go Set TyName
fvs Set TyName
seen TypeSub
sub Type TyName DefaultUni ()
b)
        TyLam ()
_ TyName
x Kind ()
k Type TyName DefaultUni ()
b
          | TyName -> Set TyName -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member TyName
x Set TyName
fvs ->
              -- This 'renameVar' makes the complexity of this function at the very least quadratic.
              ()
-> TyName
-> Kind ()
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
forall tyname (uni :: * -> *) ann.
ann
-> tyname -> Kind ann -> Type tyname uni ann -> Type tyname uni ann
TyLam () TyName
x' Kind ()
k (Type TyName DefaultUni () -> Type TyName DefaultUni ())
-> Type TyName DefaultUni () -> Type TyName DefaultUni ()
forall a b. (a -> b) -> a -> b
$ Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go (TyName -> Set TyName -> Set TyName
forall a. Ord a => a -> Set a -> Set a
Set.insert TyName
x' Set TyName
fvs) Set TyName
seen TypeSub
sub (TyName
-> TyName -> Type TyName DefaultUni () -> Type TyName DefaultUni ()
renameVar TyName
x TyName
x' Type TyName DefaultUni ()
b)
          | Bool
otherwise ->
              ()
-> TyName
-> Kind ()
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
forall tyname (uni :: * -> *) ann.
ann
-> tyname -> Kind ann -> Type tyname uni ann -> Type tyname uni ann
TyLam () TyName
x  Kind ()
k (Type TyName DefaultUni () -> Type TyName DefaultUni ())
-> Type TyName DefaultUni () -> Type TyName DefaultUni ()
forall a b. (a -> b) -> a -> b
$ Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go (TyName -> Set TyName -> Set TyName
forall a. Ord a => a -> Set a -> Set a
Set.insert TyName
x Set TyName
fvs) (TyName -> Set TyName -> Set TyName
forall a. Ord a => a -> Set a -> Set a
Set.delete TyName
x Set TyName
seen) TypeSub
sub Type TyName DefaultUni ()
b
          where x' :: TyName
x' = Set TyName -> TyName -> TyName
freshenTyNameWith (Set TyName
fvs Set TyName -> Set TyName -> Set TyName
forall a. Semigroup a => a -> a -> a
<> Getting (Set TyName) (Type TyName DefaultUni ()) TyName
-> Type TyName DefaultUni () -> Set TyName
forall a s. Getting (Set a) s a -> s -> Set a
setOf Getting (Set TyName) (Type TyName DefaultUni ()) TyName
forall tyname unique (uni :: * -> *) ann.
HasUnique tyname unique =>
Traversal' (Type tyname uni ann) tyname
Traversal' (Type TyName DefaultUni ()) TyName
ftvTy Type TyName DefaultUni ()
b) TyName
x
        TyForall ()
_ TyName
x Kind ()
k Type TyName DefaultUni ()
b
          | TyName -> Set TyName -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member TyName
x Set TyName
fvs ->
              -- This 'renameVar' makes the complexity of this function at the very least quadratic.
              ()
-> TyName
-> Kind ()
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
forall tyname (uni :: * -> *) ann.
ann
-> tyname -> Kind ann -> Type tyname uni ann -> Type tyname uni ann
TyForall () TyName
x' Kind ()
k (Type TyName DefaultUni () -> Type TyName DefaultUni ())
-> Type TyName DefaultUni () -> Type TyName DefaultUni ()
forall a b. (a -> b) -> a -> b
$ Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go (TyName -> Set TyName -> Set TyName
forall a. Ord a => a -> Set a -> Set a
Set.insert TyName
x' Set TyName
fvs) Set TyName
seen TypeSub
sub (TyName
-> TyName -> Type TyName DefaultUni () -> Type TyName DefaultUni ()
renameVar TyName
x TyName
x' Type TyName DefaultUni ()
b)
          | Bool
otherwise ->
              ()
-> TyName
-> Kind ()
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
forall tyname (uni :: * -> *) ann.
ann
-> tyname -> Kind ann -> Type tyname uni ann -> Type tyname uni ann
TyForall () TyName
x  Kind ()
k (Type TyName DefaultUni () -> Type TyName DefaultUni ())
-> Type TyName DefaultUni () -> Type TyName DefaultUni ()
forall a b. (a -> b) -> a -> b
$ Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go (TyName -> Set TyName -> Set TyName
forall a. Ord a => a -> Set a -> Set a
Set.insert TyName
x Set TyName
fvs) (TyName -> Set TyName -> Set TyName
forall a. Ord a => a -> Set a -> Set a
Set.delete TyName
x Set TyName
seen) TypeSub
sub Type TyName DefaultUni ()
b
          where x' :: TyName
x' = Set TyName -> TyName -> TyName
freshenTyNameWith (Set TyName
fvs Set TyName -> Set TyName -> Set TyName
forall a. Semigroup a => a -> a -> a
<> Getting (Set TyName) (Type TyName DefaultUni ()) TyName
-> Type TyName DefaultUni () -> Set TyName
forall a s. Getting (Set a) s a -> s -> Set a
setOf Getting (Set TyName) (Type TyName DefaultUni ()) TyName
forall tyname unique (uni :: * -> *) ann.
HasUnique tyname unique =>
Traversal' (Type tyname uni ann) tyname
Traversal' (Type TyName DefaultUni ()) TyName
ftvTy Type TyName DefaultUni ()
b) TyName
x
        TyBuiltin{} -> Type TyName DefaultUni ()
ty
        TyIFix ()
_ Type TyName DefaultUni ()
a Type TyName DefaultUni ()
b -> ()
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
forall tyname (uni :: * -> *) ann.
ann
-> Type tyname uni ann
-> Type tyname uni ann
-> Type tyname uni ann
TyIFix () (Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go Set TyName
fvs Set TyName
seen TypeSub
sub Type TyName DefaultUni ()
a) (Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go Set TyName
fvs Set TyName
seen TypeSub
sub Type TyName DefaultUni ()
b)
        TySOP ()
_ [[Type TyName DefaultUni ()]]
tyls -> () -> [[Type TyName DefaultUni ()]] -> Type TyName DefaultUni ()
forall tyname (uni :: * -> *) ann.
ann -> [[Type tyname uni ann]] -> Type tyname uni ann
TySOP () ((([Type TyName DefaultUni ()] -> [Type TyName DefaultUni ()])
-> [[Type TyName DefaultUni ()]] -> [[Type TyName DefaultUni ()]]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([Type TyName DefaultUni ()] -> [Type TyName DefaultUni ()])
 -> [[Type TyName DefaultUni ()]] -> [[Type TyName DefaultUni ()]])
-> ((Type TyName DefaultUni () -> Type TyName DefaultUni ())
    -> [Type TyName DefaultUni ()] -> [Type TyName DefaultUni ()])
-> (Type TyName DefaultUni () -> Type TyName DefaultUni ())
-> [[Type TyName DefaultUni ()]]
-> [[Type TyName DefaultUni ()]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type TyName DefaultUni () -> Type TyName DefaultUni ())
-> [Type TyName DefaultUni ()] -> [Type TyName DefaultUni ()]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (Set TyName
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
go Set TyName
fvs Set TyName
seen TypeSub
sub) [[Type TyName DefaultUni ()]]
tyls)

-- CODE REVIEW: this function is a bit strange and I don't like it. Ideas welcome for how to
-- do this better. It basically deals with the fact that we want to be careful when substituting
-- the datatypes that escape from a term into the type. It's yucky but it works.
--
-- This might not be a welcome opinion, but working with this stuff exposes some of
-- the shortcomings of the current PIR design. It would be cleaner if a PIR program was a list
-- of declarations and datatype declarations weren't in terms.
substEscape :: Set TyName
            -> TypeSub
            -> Type TyName DefaultUni ()
            -> Type TyName DefaultUni ()
substEscape :: Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
substEscape = HasCallStack =>
Bool
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
Bool
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
substTypeCustomGo Bool
True

-- See Note [Substitution in generators].
-- | Generalized substitution algorithm.
substTypeCustom
    :: HasCallStack
    => Bool
    -- ^ Nested (True) or parallel (False)
    -> TypeSub
    -> Type TyName DefaultUni ()
    -> Type TyName DefaultUni ()
substTypeCustom :: HasCallStack =>
Bool
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
substTypeCustom Bool
nested TypeSub
sub0 Type TyName DefaultUni ()
ty0 = HasCallStack =>
Bool
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
Bool
-> Set TyName
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
substTypeCustomGo Bool
nested Set TyName
fvs0 TypeSub
sub0 Type TyName DefaultUni ()
ty0 where
    fvs0 :: Set TyName
fvs0 = [Set TyName] -> Set TyName
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ([Set TyName] -> Set TyName) -> [Set TyName] -> Set TyName
forall a b. (a -> b) -> a -> b
$ TypeSub -> Set TyName
forall k a. Map k a -> Set k
Map.keysSet TypeSub
sub0 Set TyName -> [Set TyName] -> [Set TyName]
forall a. a -> [a] -> [a]
: (Type TyName DefaultUni () -> Set TyName)
-> [Type TyName DefaultUni ()] -> [Set TyName]
forall a b. (a -> b) -> [a] -> [b]
map (Getting (Set TyName) (Type TyName DefaultUni ()) TyName
-> Type TyName DefaultUni () -> Set TyName
forall a s. Getting (Set a) s a -> s -> Set a
setOf Getting (Set TyName) (Type TyName DefaultUni ()) TyName
forall tyname unique (uni :: * -> *) ann.
HasUnique tyname unique =>
Traversal' (Type tyname uni ann) tyname
Traversal' (Type TyName DefaultUni ()) TyName
ftvTy) (TypeSub -> [Type TyName DefaultUni ()]
forall k a. Map k a -> [a]
Map.elems TypeSub
sub0)

-- See Note [Substitution in generators].
-- | Regular (i.e. nested type substitution).
substType
    :: HasCallStack
    => TypeSub
    -> Type TyName DefaultUni ()
    -> Type TyName DefaultUni ()
substType :: HasCallStack =>
TypeSub -> Type TyName DefaultUni () -> Type TyName DefaultUni ()
substType = HasCallStack =>
Bool
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
Bool
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
substTypeCustom Bool
True

-- See Note [Substitution in generators].
-- | Parallel substitution.
substTypeParallel
    :: TypeSub
    -> Type TyName DefaultUni ()
    -> Type TyName DefaultUni ()
substTypeParallel :: TypeSub -> Type TyName DefaultUni () -> Type TyName DefaultUni ()
substTypeParallel = HasCallStack =>
Bool
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
Bool
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
substTypeCustom Bool
False

-- | Rename one variable to another.
renameVar :: TyName -> TyName -> Type TyName DefaultUni () -> Type TyName DefaultUni ()
renameVar :: TyName
-> TyName -> Type TyName DefaultUni () -> Type TyName DefaultUni ()
renameVar TyName
x TyName
y
    | TyName
x TyName -> TyName -> Bool
forall a. Eq a => a -> a -> Bool
== TyName
y    = Type TyName DefaultUni () -> Type TyName DefaultUni ()
forall a. a -> a
id
    | Bool
otherwise = HasCallStack =>
TypeSub -> Type TyName DefaultUni () -> Type TyName DefaultUni ()
TypeSub -> Type TyName DefaultUni () -> Type TyName DefaultUni ()
substType (TypeSub -> Type TyName DefaultUni () -> Type TyName DefaultUni ())
-> TypeSub
-> Type TyName DefaultUni ()
-> Type TyName DefaultUni ()
forall a b. (a -> b) -> a -> b
$ TyName -> Type TyName DefaultUni () -> TypeSub
forall k a. k -> a -> Map k a
Map.singleton TyName
x (() -> TyName -> Type TyName DefaultUni ()
forall tyname (uni :: * -> *) ann.
ann -> tyname -> Type tyname uni ann
TyVar () TyName
y)

-- | Find all free type variables of type `a` given substitution `sub`. If variable `x` is
-- free in `a` but in the domain of `sub` we look up `x` in `sub` and get all the free type
-- variables of the result - up to the substitution.
fvTypeR :: TypeSub -> Type TyName DefaultUni () -> Set TyName
fvTypeR :: TypeSub -> Type TyName DefaultUni () -> Set TyName
fvTypeR TypeSub
sub = Type TyName DefaultUni () -> Set TyName
go where
    go :: Type TyName DefaultUni () -> Set TyName
go = (TyName -> Set TyName) -> Set TyName -> Set TyName
forall m a. Monoid m => (a -> m) -> Set a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\TyName
v -> Set TyName
-> (Type TyName DefaultUni () -> Set TyName)
-> Maybe (Type TyName DefaultUni ())
-> Set TyName
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (TyName -> Set TyName
forall a. a -> Set a
Set.singleton TyName
v) Type TyName DefaultUni () -> Set TyName
go (Maybe (Type TyName DefaultUni ()) -> Set TyName)
-> Maybe (Type TyName DefaultUni ()) -> Set TyName
forall a b. (a -> b) -> a -> b
$ TyName -> TypeSub -> Maybe (Type TyName DefaultUni ())
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TyName
v TypeSub
sub) (Set TyName -> Set TyName)
-> (Type TyName DefaultUni () -> Set TyName)
-> Type TyName DefaultUni ()
-> Set TyName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting (Set TyName) (Type TyName DefaultUni ()) TyName
-> Type TyName DefaultUni () -> Set TyName
forall a s. Getting (Set a) s a -> s -> Set a
setOf Getting (Set TyName) (Type TyName DefaultUni ()) TyName
forall tyname unique (uni :: * -> *) ann.
HasUnique tyname unique =>
Traversal' (Type tyname uni ann) tyname
Traversal' (Type TyName DefaultUni ()) TyName
ftvTy

-- * Generators for substitutions

-- | Get the free type variables in a type along with how many
-- times they occur. The elements of the map are guaranteed to be
-- non-zero.
fvTypeBag :: Type TyName DefaultUni () -> Map TyName Int
fvTypeBag :: Type TyName DefaultUni () -> Map TyName Int
fvTypeBag = MultiSet TyName -> Map TyName Int
forall a. MultiSet a -> Map a Int
toMap (MultiSet TyName -> Map TyName Int)
-> (Type TyName DefaultUni () -> MultiSet TyName)
-> Type TyName DefaultUni ()
-> Map TyName Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting (MultiSet TyName) (Type TyName DefaultUni ()) TyName
-> Type TyName DefaultUni () -> MultiSet TyName
forall a s. Getting (MultiSet a) s a -> s -> MultiSet a
multiSetOf Getting (MultiSet TyName) (Type TyName DefaultUni ()) TyName
forall tyname unique (uni :: * -> *) ann.
HasUnique tyname unique =>
Traversal' (Type tyname uni ann) tyname
Traversal' (Type TyName DefaultUni ()) TyName
ftvTy

-- | Generate a type substitution mapping some of the variables in the given context to some
-- arbitrary types (valid in the context).
genSubst :: TypeCtx -> Gen TypeSub
genSubst :: TypeCtx -> Gen TypeSub
genSubst TypeCtx
ctx0 = do
  [(TyName, Kind ())]
xks <- [(TyName, Kind ())] -> Gen [(TyName, Kind ())]
forall a. [a] -> Gen [a]
sublistOf ([(TyName, Kind ())] -> Gen [(TyName, Kind ())])
-> ([(TyName, Kind ())] -> Gen [(TyName, Kind ())])
-> [(TyName, Kind ())]
-> Gen [(TyName, Kind ())]
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< [(TyName, Kind ())] -> Gen [(TyName, Kind ())]
forall a. [a] -> Gen [a]
shuffle ([(TyName, Kind ())] -> Gen [(TyName, Kind ())])
-> [(TyName, Kind ())] -> Gen [(TyName, Kind ())]
forall a b. (a -> b) -> a -> b
$ TypeCtx -> [(TyName, Kind ())]
forall k a. Map k a -> [(k, a)]
Map.toList TypeCtx
ctx0
  TypeCtx
-> TypeSub -> Map TyName Int -> [(TyName, Kind ())] -> Gen TypeSub
go TypeCtx
ctx0 TypeSub
forall k a. Map k a
Map.empty Map TyName Int
forall k a. Map k a
Map.empty [(TyName, Kind ())]
xks
  where
    -- Counts is used to balance the ratio between the number of times a variable @x@ occurs in the
    -- substitution and the size of the type it maps to - the more times @x@ occurs the smaller the
    -- type it maps to needs to be to avoid blowup.
    go :: TypeCtx
-> TypeSub -> Map TyName Int -> [(TyName, Kind ())] -> Gen TypeSub
go TypeCtx
_   TypeSub
sub Map TyName Int
_      []            = TypeSub -> Gen TypeSub
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeSub
sub
    go TypeCtx
ctx TypeSub
sub Map TyName Int
counts ((TyName
x, Kind ()
k) : [(TyName, Kind ())]
xs) = do
      let
          -- @x@ is taken out from the context, because we're going to map it to a type valid in the
          -- context without @x@.
          ctx' :: TypeCtx
ctx' = TyName -> TypeCtx -> TypeCtx
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete TyName
x TypeCtx
ctx
          -- How many times @x@ occurs in all the so far generated types (the ones that are in the
          -- codomain of @sub@).
          w :: Int
w = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
1 (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ TyName -> Map TyName Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TyName
x Map TyName Int
counts
      Type TyName DefaultUni ()
ty <- (Int -> Gen (Type TyName DefaultUni ()))
-> Gen (Type TyName DefaultUni ())
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen (Type TyName DefaultUni ()))
 -> Gen (Type TyName DefaultUni ()))
-> (Int -> Gen (Type TyName DefaultUni ()))
-> Gen (Type TyName DefaultUni ())
forall a b. (a -> b) -> a -> b
$ \ Int
n -> Int
-> Gen (Type TyName DefaultUni ())
-> Gen (Type TyName DefaultUni ())
forall a. HasCallStack => Int -> Gen a -> Gen a
resize (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
w) (Gen (Type TyName DefaultUni ())
 -> Gen (Type TyName DefaultUni ()))
-> Gen (Type TyName DefaultUni ())
-> Gen (Type TyName DefaultUni ())
forall a b. (a -> b) -> a -> b
$ TypeCtx -> Kind () -> Gen (Type TyName DefaultUni ())
genTypeWithCtx TypeCtx
ctx' Kind ()
k
      let -- Scale occurrences of all free variables of @ty@ according to how many times @x@
          -- (the variables that is being substituted for) occurs in the so far generated
          -- substitution.
          moreCounts :: Map TyName Int
moreCounts = (Int -> Int) -> Map TyName Int -> Map TyName Int
forall a b. (a -> b) -> Map TyName a -> Map TyName b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
w) (Map TyName Int -> Map TyName Int)
-> Map TyName Int -> Map TyName Int
forall a b. (a -> b) -> a -> b
$ Type TyName DefaultUni () -> Map TyName Int
fvTypeBag Type TyName DefaultUni ()
ty
          sub' :: TypeSub
sub'       = TyName -> Type TyName DefaultUni () -> TypeSub -> TypeSub
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert TyName
x Type TyName DefaultUni ()
ty TypeSub
sub
          counts' :: Map TyName Int
counts'    = (Int -> Int -> Int)
-> Map TyName Int -> Map TyName Int -> Map TyName Int
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Map TyName Int
counts Map TyName Int
moreCounts
      TypeCtx
-> TypeSub -> Map TyName Int -> [(TyName, Kind ())] -> Gen TypeSub
go TypeCtx
ctx' TypeSub
sub' Map TyName Int
counts' [(TyName, Kind ())]
xs

shrinkSubst :: TypeCtx -> TypeSub -> [TypeSub]
shrinkSubst :: TypeCtx -> TypeSub -> [TypeSub]
shrinkSubst TypeCtx
ctx0 = ([(TyName, Type TyName DefaultUni ())] -> TypeSub)
-> [[(TyName, Type TyName DefaultUni ())]] -> [TypeSub]
forall a b. (a -> b) -> [a] -> [b]
map [(TyName, Type TyName DefaultUni ())] -> TypeSub
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([[(TyName, Type TyName DefaultUni ())]] -> [TypeSub])
-> (TypeSub -> [[(TyName, Type TyName DefaultUni ())]])
-> TypeSub
-> [TypeSub]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((TyName, Type TyName DefaultUni ())
 -> [(TyName, Type TyName DefaultUni ())])
-> [(TyName, Type TyName DefaultUni ())]
-> [[(TyName, Type TyName DefaultUni ())]]
forall a. (a -> [a]) -> [a] -> [[a]]
forall (f :: * -> *) a. Arbitrary1 f => (a -> [a]) -> f a -> [f a]
liftShrink (TyName, Type TyName DefaultUni ())
-> [(TyName, Type TyName DefaultUni ())]
shrinkTy ([(TyName, Type TyName DefaultUni ())]
 -> [[(TyName, Type TyName DefaultUni ())]])
-> (TypeSub -> [(TyName, Type TyName DefaultUni ())])
-> TypeSub
-> [[(TyName, Type TyName DefaultUni ())]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSub -> [(TyName, Type TyName DefaultUni ())]
forall k a. Map k a -> [(k, a)]
Map.toList
  where
    shrinkTy :: (TyName, Type TyName DefaultUni ())
-> [(TyName, Type TyName DefaultUni ())]
shrinkTy (TyName
x, Type TyName DefaultUni ()
ty) = (,) TyName
x (Type TyName DefaultUni () -> (TyName, Type TyName DefaultUni ()))
-> [Type TyName DefaultUni ()]
-> [(TyName, Type TyName DefaultUni ())]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HasCallStack =>
TypeCtx
-> Kind ()
-> Type TyName DefaultUni ()
-> [Type TyName DefaultUni ()]
TypeCtx
-> Kind ()
-> Type TyName DefaultUni ()
-> [Type TyName DefaultUni ()]
shrinkTypeAtKind (TypeCtx -> Type TyName DefaultUni () -> TypeCtx
forall {unique} {k} {a} {uni :: * -> *} {ann}.
(Coercible unique Int, Ord k, HasUnique k unique) =>
Map k a -> Type k uni ann -> Map k a
pruneCtx TypeCtx
ctx0 Type TyName DefaultUni ()
ty) Kind ()
k Type TyName DefaultUni ()
ty
      where k :: Kind ()
k = Kind () -> Maybe (Kind ()) -> Kind ()
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> Kind ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> Kind ()) -> [Char] -> Kind ()
forall a b. (a -> b) -> a -> b
$ [Char]
"internal error: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TyName -> [Char]
forall a. Show a => a -> [Char]
show TyName
x [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" not found") (Maybe (Kind ()) -> Kind ()) -> Maybe (Kind ()) -> Kind ()
forall a b. (a -> b) -> a -> b
$ TyName -> TypeCtx -> Maybe (Kind ())
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TyName
x TypeCtx
ctx0
    pruneCtx :: Map k a -> Type k uni ann -> Map k a
pruneCtx Map k a
ctx Type k uni ann
ty = Map k a
ctx Map k a -> Set k -> Map k a
forall k a. Ord k => Map k a -> Set k -> Map k a
`Map.restrictKeys` Getting (Set k) (Type k uni ann) k -> Type k uni ann -> Set k
forall a s. Getting (Set a) s a -> s -> Set a
setOf Getting (Set k) (Type k uni ann) k
forall tyname unique (uni :: * -> *) ann.
HasUnique tyname unique =>
Traversal' (Type tyname uni ann) tyname
Traversal' (Type k uni ann) k
ftvTy Type k uni ann
ty