{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-orphans #-}

{-|
Module      : Polytropos.QuickCheck
Description : QuickCheck generators and slot-indexed shrinkers for Polytropos records and variants
Copyright   : (c) 2026 The Polytropos Authors
License     : EUPL-1.2 OR Apache-2.0
Maintainer  : nadia.yvette.chambers@ik.me
Stability   : experimental
Portability : non-portable (GHC extensions)

Integrates @polytropos@ open records and variants with QuickCheck
property-based testing.

= Architectural Overview & Rationale

== Records

'Arbitrary (Record r)' generates each field value via that field's own
'Arbitrary' instance (the walk is heterogeneous: rows may freely mix field
types) and packs the results straight into the contiguous 'Data.Primitive.SmallArray.SmallArray'
payload — with no intermediate tuples or generic-representation trees.

Shrinking is slot-indexed: each field slot is shrunk independently, one
field at a time, so shrinking a \(k\)-field record yields \(O(k)\) candidates
per level instead of the full Cartesian product over per-field shrink lists.

== Variants

'Arbitrary (Variant s)' uniformly selects a constructor tag index and
generates that constructor's payload fields. Shrinking is
branch-preserving: only the payload of the currently active constructor
is shrunk; the tag never changes, so a counterexample stays on the
constructor branch that triggered the property failure.

= When to Use This Module

* When property-testing systems that use @polytropos@ open records and variants.

> {-# LANGUAGE DataKinds #-}
> {-# LANGUAGE OverloadedRecordDot #-}
> {-# LANGUAGE TypeApplications #-}
>
> import Test.QuickCheck
> import Polytropos
> import Polytropos.QuickCheck
>
> type User = Record '["id" ':= Int, "active" ':= Bool]
>
> prop_userPositiveId :: Property
> prop_userPositiveId = forAll (arbitrary @User) $ \u ->
>   u.id == u.id
-}
module Polytropos.QuickCheck
  ( -- * Row and Sum Constraint Walks
    AllArbitrary (..)
  , AllCtorsArbitrary (..)
  ) where

import Data.Functor.Identity (Identity (..))
import Data.Proxy            (Proxy (..))
import GHC.Exts              (Any, toList)
import Polytropos            (Record, Variant)
import Polytropos.Core       (Pair (..), Row, RowRep (..), SumRow, noInlineUnsafeCo,
                              injectRep, productRep, repTag, repVals)
import Test.QuickCheck       (Arbitrary (..), Gen, choose)

{-------------------------------------------------------------------------------
  Row walk
-------------------------------------------------------------------------------}

-- | Heterogeneous generation walk: every field of the row @r@ has an
-- 'Arbitrary' instance at its own type.
class AllArbitrary (r :: Row) where
  -- | Generate a value for each payload slot, in row order.
  genFields :: Proxy r -> Gen [Identity Any]
  -- | Slot-indexed shrinks of the payload: one field shrinks at a time.
  shrinkFields :: Proxy r -> [Identity Any] -> [[Identity Any]]

instance AllArbitrary '[] where
  genFields :: Proxy '[] -> Gen [Identity Any]
genFields Proxy '[]
_  = [Identity Any] -> Gen [Identity Any]
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
  shrinkFields :: Proxy '[] -> [Identity Any] -> [[Identity Any]]
shrinkFields Proxy '[]
_ [Identity Any]
_ = []

instance (Arbitrary a, AllArbitrary r) => AllArbitrary ((n ':= a) ': r) where
  genFields :: Proxy ((n ':= a) : r) -> Gen [Identity Any]
genFields Proxy ((n ':= a) : r)
_ = do
    v <- Gen a
forall a. Arbitrary a => Gen a
arbitrary :: Gen a
    vs <- genFields (Proxy @r)
    pure (Identity (noInlineUnsafeCo v :: Any) : vs)

  shrinkFields :: Proxy ((n ':= a) : r) -> [Identity Any] -> [[Identity Any]]
shrinkFields Proxy ((n ':= a) : r)
_ (Identity Any
x : [Identity Any]
xs) =
    [ Any -> Identity Any
forall a. a -> Identity a
Identity (a -> Any
forall a b. a -> b
noInlineUnsafeCo a
v :: Any) Identity Any -> [Identity Any] -> [Identity Any]
forall a. a -> [a] -> [a]
: [Identity Any]
xs
    | a
v <- a -> [a]
forall a. Arbitrary a => a -> [a]
shrink (Any -> a
forall a b. a -> b
noInlineUnsafeCo Any
x :: a) ]
    [[Identity Any]] -> [[Identity Any]] -> [[Identity Any]]
forall a. [a] -> [a] -> [a]
++ [ Any -> Identity Any
forall a. a -> Identity a
Identity Any
x Identity Any -> [Identity Any] -> [Identity Any]
forall a. a -> [a] -> [a]
: [Identity Any]
rest | [Identity Any]
rest <- Proxy r -> [Identity Any] -> [[Identity Any]]
forall (r :: Row).
AllArbitrary r =>
Proxy r -> [Identity Any] -> [[Identity Any]]
shrinkFields (forall (t :: Row). Proxy t
forall {k} (t :: k). Proxy t
Proxy @r) [Identity Any]
xs ]
  shrinkFields Proxy ((n ':= a) : r)
_ [] = []

{-------------------------------------------------------------------------------
  Sum walk
-------------------------------------------------------------------------------}

-- | Heterogeneous generation walk over a sum row: every constructor's
-- field row satisfies 'AllArbitrary'.
class AllCtorsArbitrary (s :: SumRow) where
  -- | Number of constructors in the sum.
  numCtors :: Proxy s -> Int
  -- | Generate the payload fields of constructor @i@ (0-based).
  genCtorFields :: Proxy s -> Int -> Gen [Identity Any]
  -- | Slot-indexed shrinks of the payload of constructor @i@.
  shrinkCtorFields :: Proxy s -> Int -> [Identity Any] -> [[Identity Any]]

instance AllCtorsArbitrary '[] where
  numCtors :: Proxy '[] -> Int
numCtors Proxy '[]
_          = Int
0
  genCtorFields :: Proxy '[] -> Int -> Gen [Identity Any]
genCtorFields Proxy '[]
_ Int
_   = [Identity Any] -> Gen [Identity Any]
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
  shrinkCtorFields :: Proxy '[] -> Int -> [Identity Any] -> [[Identity Any]]
shrinkCtorFields Proxy '[]
_ Int
_ [Identity Any]
_ = []

instance (AllArbitrary fs, AllCtorsArbitrary s)
      => AllCtorsArbitrary ((c ':= fs) ': s) where
  numCtors :: Proxy ((c ':= fs) : s) -> Int
numCtors Proxy ((c ':= fs) : s)
_ = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Proxy s -> Int
forall (s :: SumRow). AllCtorsArbitrary s => Proxy s -> Int
numCtors (forall (t :: SumRow). Proxy t
forall {k} (t :: k). Proxy t
Proxy @s)

  genCtorFields :: Proxy ((c ':= fs) : s) -> Int -> Gen [Identity Any]
genCtorFields Proxy ((c ':= fs) : s)
_ Int
0 = Proxy fs -> Gen [Identity Any]
forall (r :: Row). AllArbitrary r => Proxy r -> Gen [Identity Any]
genFields (forall (t :: Row). Proxy t
forall {k} (t :: k). Proxy t
Proxy @fs)
  genCtorFields Proxy ((c ':= fs) : s)
_ Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Char] -> Gen [Identity Any]
forall a. HasCallStack => [Char] -> a
error [Char]
"polytropos-quickcheck: negative constructor index"
                    | Bool
otherwise = Proxy s -> Int -> Gen [Identity Any]
forall (s :: SumRow).
AllCtorsArbitrary s =>
Proxy s -> Int -> Gen [Identity Any]
genCtorFields (forall (t :: SumRow). Proxy t
forall {k} (t :: k). Proxy t
Proxy @s) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

  shrinkCtorFields :: Proxy ((c ':= fs) : s) -> Int -> [Identity Any] -> [[Identity Any]]
shrinkCtorFields Proxy ((c ':= fs) : s)
_ Int
0 [Identity Any]
xs = Proxy fs -> [Identity Any] -> [[Identity Any]]
forall (r :: Row).
AllArbitrary r =>
Proxy r -> [Identity Any] -> [[Identity Any]]
shrinkFields (forall (t :: Row). Proxy t
forall {k} (t :: k). Proxy t
Proxy @fs) [Identity Any]
xs
  shrinkCtorFields Proxy ((c ':= fs) : s)
_ Int
i [Identity Any]
xs = Proxy s -> Int -> [Identity Any] -> [[Identity Any]]
forall (s :: SumRow).
AllCtorsArbitrary s =>
Proxy s -> Int -> [Identity Any] -> [[Identity Any]]
shrinkCtorFields (forall (t :: SumRow). Proxy t
forall {k} (t :: k). Proxy t
Proxy @s) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Identity Any]
xs

{-------------------------------------------------------------------------------
  QuickCheck instances
-------------------------------------------------------------------------------}

instance AllArbitrary r => Arbitrary (Record r) where
  arbitrary :: Gen (Record r)
arbitrary = Rep Identity () -> Record r
forall {k} k1 (f :: * -> *) (r :: RowType) (row :: k1) (a :: k).
Rep f a -> RowRep f r row a
RowRep (Rep Identity () -> Record r)
-> ([Identity Any] -> Rep Identity ())
-> [Identity Any]
-> Record r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Identity Any] -> Rep Identity ()
forall {k} (f :: * -> *) (a :: k). [f Any] -> Rep f a
productRep ([Identity Any] -> Record r)
-> Gen [Identity Any] -> Gen (Record r)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Proxy r -> Gen [Identity Any]
forall (r :: Row). AllArbitrary r => Proxy r -> Gen [Identity Any]
genFields (forall (t :: Row). Proxy t
forall {k} (t :: k). Proxy t
Proxy @r)
  shrink :: Record r -> [Record r]
shrink (RowRep Rep Identity ()
rep) =
    [ Rep Identity () -> Record r
forall {k} k1 (f :: * -> *) (r :: RowType) (row :: k1) (a :: k).
Rep f a -> RowRep f r row a
RowRep ([Identity Any] -> Rep Identity ()
forall {k} (f :: * -> *) (a :: k). [f Any] -> Rep f a
productRep [Identity Any]
vs)
    | [Identity Any]
vs <- Proxy r -> [Identity Any] -> [[Identity Any]]
forall (r :: Row).
AllArbitrary r =>
Proxy r -> [Identity Any] -> [[Identity Any]]
shrinkFields (forall (t :: Row). Proxy t
forall {k} (t :: k). Proxy t
Proxy @r) (SmallArray (Identity Any) -> [Item (SmallArray (Identity Any))]
forall l. IsList l => l -> [Item l]
toList (Rep Identity () -> SmallArray (Identity Any)
forall {k} (f :: * -> *) (a :: k). Rep f a -> SmallArray (f Any)
repVals Rep Identity ()
rep)) ]

instance AllCtorsArbitrary s => Arbitrary (Variant s) where
  arbitrary :: Gen (Variant s)
arbitrary = do
    let n :: Int
n = Proxy s -> Int
forall (s :: SumRow). AllCtorsArbitrary s => Proxy s -> Int
numCtors (forall (t :: SumRow). Proxy t
forall {k} (t :: k). Proxy t
Proxy @s)
    if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (Int
0 :: Int)
      then Variant s -> Gen (Variant s)
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Rep Identity () -> Variant s
forall {k} k1 (f :: * -> *) (r :: RowType) (row :: k1) (a :: k).
Rep f a -> RowRep f r row a
RowRep (Word -> [Identity Any] -> Rep Identity ()
forall {k} (f :: * -> *) (a :: k). Word -> [f Any] -> Rep f a
injectRep Word
0 []))
      else do
        i <- (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
0, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        vs <- genCtorFields (Proxy @s) i
        pure (RowRep (injectRep (fromIntegral (i :: Int)) vs))
  shrink :: Variant s -> [Variant s]
shrink (RowRep Rep Identity ()
rep) =
    let tag :: Word
tag = Rep Identity () -> Word
forall {k} (f :: * -> *) (a :: k). Rep f a -> Word
repTag Rep Identity ()
rep
    in [ Rep Identity () -> Variant s
forall {k} k1 (f :: * -> *) (r :: RowType) (row :: k1) (a :: k).
Rep f a -> RowRep f r row a
RowRep (Word -> [Identity Any] -> Rep Identity ()
forall {k} (f :: * -> *) (a :: k). Word -> [f Any] -> Rep f a
injectRep Word
tag [Identity Any]
vs)
       | [Identity Any]
vs <- Proxy s -> Int -> [Identity Any] -> [[Identity Any]]
forall (s :: SumRow).
AllCtorsArbitrary s =>
Proxy s -> Int -> [Identity Any] -> [[Identity Any]]
shrinkCtorFields (forall (t :: SumRow). Proxy t
forall {k} (t :: k). Proxy t
Proxy @s) (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
tag)
                                            (SmallArray (Identity Any) -> [Item (SmallArray (Identity Any))]
forall l. IsList l => l -> [Item l]
toList (Rep Identity () -> SmallArray (Identity Any)
forall {k} (f :: * -> *) (a :: k). Rep f a -> SmallArray (f Any)
repVals Rep Identity ()
rep)) ]