Gen.hs

{-# LANGUAGE TupleSections #-}

module HaskellWorks.Data.BalancedParens.Gen
  ( BP(..)
  , count
  , bpBools
  , showBps
  , storableVector
  , bpParensSeq
  , vector
  , vec2
  , randomRmm
  , randomRmm2
  ) where

import Data.Coerce
import Data.Semigroup                             ((<>))
import Data.Word
import HaskellWorks.Data.BalancedParens.ParensSeq (ParensSeq)
import HaskellWorks.Data.Positioning
import Hedgehog

import qualified Data.Vector                                   as DV
import qualified Data.Vector.Storable                          as DVS
import qualified HaskellWorks.Data.BalancedParens.ParensSeq    as PS
import qualified HaskellWorks.Data.BalancedParens.RangeMinMax  as RMM
import qualified HaskellWorks.Data.BalancedParens.RangeMinMax2 as RMM2
import qualified Hedgehog.Gen                                  as G
import qualified Hedgehog.Range                                as R

count :: MonadGen m => Range Count -> m Count
count r = coerce <$> G.word64 (coerce <$> r)

data LR a = L a Int | R a Int deriving (Eq, Show)

newtype BP = BP [Bool] deriving Eq

showBps :: [Bool] -> String
showBps = fmap fromBool
  where fromBool True  = '('
        fromBool False = ')'

bpBools' :: MonadGen m => Int -> (Int, [Bool], [Bool], Int) -> m [Bool]
bpBools' n (ln, lt, rt, rn) = if n <= 0
  then return (reverse lt <> rt)
  else if ln - rn >= n
    then return (reverse lt <> replicate n False <> rt)
    else if rn - ln >= n
      then return (reverse lt <> replicate n True <> rt)
      else do
        decision <- case (ln, rn) of
          (0, 0) -> G.element [L '(' 1,                         R ')' 1]
          (0, _) -> G.element [L '(' 1,             R '(' (-1), R ')' 1]
          (_, 0) -> G.element [L '(' 1, L ')' (-1),             R ')' 1]
          _      -> G.element [L '(' 1,                         R ')' 1]

        case decision of
          L p d -> bpBools' (n - 1) (ln + d, toBool p:lt,          rt, rn    )
          R p d -> bpBools' (n - 1) (ln    ,          lt, toBool p:rt, rn + d)
  where toBool '(' = True
        toBool  _  = False

bpBools ::  MonadGen m => Range Int -> m [Bool]
bpBools r = do
  n <- G.int r
  bpBools' (n * 2) (0, [], [], 0)

bpParensSeq ::  MonadGen m => Range Int -> m ParensSeq
bpParensSeq = fmap PS.fromBools . bpBools

storableVector :: (MonadGen m, DVS.Storable a) => Range Int -> m a -> m (DVS.Vector a)
storableVector r g = DVS.fromList <$> G.list r g

vector :: MonadGen m => Range Int -> m a -> m (DV.Vector a)
vector r g = DV.fromList <$> G.list r g

vec2 :: MonadGen m => m a -> m (a, a)
vec2 g = (,) <$> g <*> g

randomRmm :: MonadGen m => Range Int -> m (RMM.RangeMinMax (DVS.Vector Word64))
randomRmm r = do
  v <- storableVector (fmap (64 *) r) (G.word64 R.constantBounded)
  return (RMM.mkRangeMinMax v)

randomRmm2 :: MonadGen m => Range Int -> m (RMM2.RangeMinMax2 (DVS.Vector Word64))
randomRmm2 r = do
  v <- storableVector (fmap (64 *) r) (G.word64 R.constantBounded)
  return (RMM2.mkRangeMinMax2 v)