0 | module Hedgehog.Internal.Function

  2 | import Data.Colist

  3 | import Data.Cotree

  4 | import Data.Either

  5 | import Data.Nat

  6 | import Data.String

  8 | import Hedgehog.Internal.Gen

  9 | import Hedgehog.Internal.Range

 10 | import Hedgehog.Internal.Util

 12 | import System.Random.Pure.StdGen

 14 | %default total

 34 | public export

 35 | interface Cogen a where

 36 |   constructor MkCogen

 37 |   perturb : a -> StdGen -> StdGen

 68 | shiftArg : StdGen -> StdGen

 69 | shiftArg = variant 33 . snd . split . variant 31

 77 | cogen : Cogen a => a -> Gen b -> Gen b

 78 | cogen x g = MkGen $ \sz, sd => unGen g sz $ perturb x sd

 81 | Cogen Unit where

 82 |   perturb _ = id

 85 | Cogen (Equal x y) where

 86 |   perturb Refl = id

 89 | Cogen Bool where

 90 |   perturb True  = variant 0

 91 |   perturb False = variant 1

 94 | Cogen Nat where

 95 |   perturb = variant

 98 | Cogen Integer where

 99 |   perturb = variant . cast

102 | Cogen Bits64 where

103 |   perturb = variant . cast

106 | Cogen Bits16 where perturb = variant . cast

109 | Cogen Bits8 where perturb = variant . cast

112 | Cogen Int64 where perturb = variant . cast

115 | Cogen Int16 where perturb = variant . cast

118 | Cogen Int8 where perturb = variant . cast

121 | Cogen Int where perturb = variant . cast

124 | Cogen Char where perturb = variant . cast

127 | Cogen Void where

128 |   perturb _ _ impossible

131 | Cogen a => Cogen b => Cogen (a, b) where

132 |   perturb (x, y) = perturb x . shiftArg . perturb y . shiftArg

135 | Cogen a => Cogen b => Cogen (Either a b) where

136 |   perturb $ Left  x = perturb x . shiftArg . variant 0

137 |   perturb $ Right y = perturb y . shiftArg . variant 1

140 | Cogen a => Cogen (Maybe a) where

141 |   perturb Nothing  = variant 0

142 |   perturb (Just x) = perturb x . shiftArg . variant 1

145 | Cogen a => Cogen (List a) where

146 |   perturb []      = variant 0

147 |   perturb (x::xs) = perturb xs . shiftArg . perturb x . shiftArg . variant 1

150 | Cogen String where

151 |   perturb = perturb . fastUnpack

164 | function_ : Cogen a => Gen b -> Gen (a -> b)

165 | function_ bg =

166 |   MkGen $ \sz, sd => singleton $ \x => value $ unGen bg sz $ perturb x sd

178 | depfun_ :

179 |      {auto _ : Cogen a}

180 |   -> {0 b : a -> Type}

181 |   -> ((x : a) -> Gen $ b x)

182 |   -> Gen ((x : a) -> b x)

183 | depfun_ bg =

184 |   MkGen $ \sz, sd => singleton $ \x => value $ unGen (bg x) sz $ perturb x sd

196 | dargfun_ :

197 |      {0 b    : a -> Type}

198 |   -> {auto _ : {0 x : a} -> Cogen (b x)}

199 |   -> Gen c

200 |   -> Gen ({0 x : a} -> b x -> c)

201 | dargfun_ bg =

202 |   MkGen $ \sz, sd => singleton $ \x => value $ unGen bg sz $ perturb x sd

214 | dargdepfun_ :

215 |      {0 b : a -> Type}

216 |   -> {0 c : {0 x : a} -> b x -> Type}

217 |   -> {auto _ : {0 x : a} -> Cogen (b x)}

218 |   -> ({0 x : a} -> (y : b x) -> Gen (c y))

219 |   -> Gen ({0 x : a} -> (y : b x) -> c y)

220 | dargdepfun_ bg =

221 |   MkGen $ \sz, sd => singleton $ \x => value $ unGen (bg x) sz $ perturb x sd

230 | export infixr 5 :->

240 | public export

241 | data (:->) : Type -> Type -> Type where

242 |   FUnit : c -> () :-> c

243 |   FNil  : a :-> c

244 |   FPair : Lazy (a :-> b :-> c) -> (a, b) :-> c

245 |   FSum  : Lazy (a :-> c) -> Lazy (b :-> c) -> Either a b :-> c

246 |   FMap  : (a -> b) -> (b -> a) -> Lazy (b :-> c) -> a :-> c

249 | Functor ((:->) a) where

250 |   map f $ FUnit c    = FUnit $ f c

251 |   map _ $ FNil       = FNil

252 |   map f $ FPair a    = FPair $ map (assert_total $ map f) a

253 |   map f $ FSum a b   = FSum (map f a) (map f b)

254 |   map f $ FMap a b c = FMap a b $ map f c

256 | table : a :-> c -> List (a, c)

257 | table (FUnit c) = [((), c)]

258 | table FNil      = []

259 | table (FPair f) = do

260 |   (a, bc) <- table f

261 |   (b, c) <- assert_total table bc

262 |   pure ((a, b), c)

263 | table (FSum a b) =

264 |      [(Left x, c)  | (x, c) <- table a]

265 |   ++ [(Right x, c) | (x, c) <- table b]

266 | table (FMap _ g a) = mapFst g <$> table a

268 | public export

269 | interface Cogen a => ShrCogen a where

270 |   constructor MkShrCogen

271 |   build : (a -> c) -> a :-> c

274 | ShrCogen Void where

275 |   build _ = FNil

278 | ShrCogen Unit where

279 |   build f = FUnit $ f ()

282 | ShrCogen a => ShrCogen b => ShrCogen (a, b) where

283 |   build f = FPair $ build $ \a => build $ \b => f (a, b)

286 | ShrCogen a => ShrCogen b => ShrCogen (Either a b) where

287 |   build f = FSum (build $ f . Left) (build $ f . Right)

295 | via : ShrCogen b => (a -> b) -> (b -> a) -> (a -> c) -> a :-> c

296 | via a b f = FMap a b $ build $ f . b

299 | ShrCogen a => ShrCogen (Maybe a) where

300 |   build = via (maybeToEither ()) eitherToMaybe

303 | ShrCogen Bool where

304 |   build = via toEither fromEither

306 |     where

307 |       toEither : Bool -> Either Unit Unit

308 |       toEither True  = Left ()

309 |       toEither False = Right ()

310 |       fromEither : Either Unit Unit -> Bool

311 |       fromEither $ Left ()  = True

312 |       fromEither $ Right () = False

315 | ShrCogen a => ShrCogen (List a) where

316 |   build = assert_total via toEither fromEither

318 |     where

319 |       toEither : List a -> Either Unit (a, List a)

320 |       toEither []      = Left ()

321 |       toEither (x::xs) = Right (x, xs)

322 |       fromEither : Either Unit (a, List a) -> List a

323 |       fromEither (Left ())       = []

324 |       fromEither (Right (x, xs)) = x::xs

327 | ShrCogen (Equal x x) where

328 |   build = via (const ()) (const Refl)

331 | ShrCogen Integer where

332 |   build = via toBits fromBits

334 |     where

335 |       toBits : Integer -> (Bool, List Bool)

336 |       toBits n = if n >= 0 then (True, go [] n) else (False, go [] $ -n - 1)

338 |         where

339 |           go : List Bool -> Integer -> List Bool

340 |           go bits x =

341 |             if x == 0

342 |               then bits

343 |               else go ((mod x 2 == 1) :: bits) (assert_smaller x $ div x 2)

345 |       fromBits : (Bool, List Bool) -> Integer

346 |       fromBits (sign, bits) = do

347 |         let body = foldl (\acc, b => acc * the Integer 2 + if b then 1 else 0) 0 bits

348 |         if sign then body else negate $ body + 1

351 | ShrCogen Nat where build = via {b=Integer} cast cast

354 | ShrCogen Int where build = via {b=Integer} cast cast

357 | ShrCogen Int8 where build = via {b=Integer} cast cast

360 | ShrCogen Int16 where build = via {b=Integer} cast cast

363 | ShrCogen Int64 where build = via {b=Integer} cast cast

366 | ShrCogen Bits8 where build = via {b=Nat} cast cast

369 | ShrCogen Bits16 where build = via {b=Nat} cast cast

372 | ShrCogen Bits64 where build = via {b=Nat} cast cast

375 | ShrCogen Char where build = via {b=Nat} cast cast

378 | ShrCogen String where build = via fastUnpack fastPack

380 | apply' : a :-> b -> a -> Maybe b

381 | apply' (FUnit c)    ()        = Just c

382 | apply' FNil         _         = Nothing

383 | apply' (FPair f)    (a, b)    = assert_total $ apply' !(apply' f a) b

384 | apply' (FSum f _)   (Left a)  = apply' f a

385 | apply' (FSum _ g)   (Right a) = apply' g a

386 | apply' (FMap f _ g) a         = apply' g (f a)

388 | shrinkFn : (b -> Inf (Colist b)) -> a :-> b -> Colist $ a :-> b

389 | shrinkFn shr (FUnit a)  = FUnit <$> shr a

390 | shrinkFn _   FNil       = []

391 | shrinkFn shr (FPair f)  =

392 |   shrinkFn (delay . assert_total (shrinkFn shr)) f <&>

393 |     \case FNil => FNil; a => FPair a

394 | shrinkFn shr (FSum a b) =

395 |   map (\case FSum FNil FNil => FNil; x => x) $

396 |     (if notFNil b then [ FSum a FNil ] else []) ++

397 |     (if notFNil a then [ FSum FNil b ] else []) ++

398 |     map (`FSum` b) (shrinkFn shr a) ++

399 |     map (a `FSum`) (shrinkFn shr b)

401 |   where

402 |     (++) : forall a. Colist a -> Inf (Colist a) -> Colist a

403 |     []      ++ ys = ys

404 |     (x::xs) ++ ys = x :: (xs ++ ys)

406 |     notFNil : forall a, b. a :-> b -> Bool

407 |     notFNil FNil = False

408 |     notFNil _   = True

409 | shrinkFn shr (FMap f g a) =

410 |   shrinkFn shr a <&> \case FNil => FNil; x => FMap f g x

414 | data Fn a b = MkFn b (a :-> Cotree b)

417 | Show a => Show b => Show (Fn a b) where

418 |   show (MkFn xb xa) = unlines $ (table xa <&> showCase) ++ ["_ -> " ++ show xb]

420 |     where

421 |       showCase : (a, Cotree b) -> String

422 |       showCase (lhs, rhs) = show lhs ++ " -> " ++ show rhs.value

433 | function : ShrCogen a => Gen b -> Gen (Fn a b)

434 | function gb = [| MkFn gb (genFn $ \a => cogen a gb) |]

436 |   where

437 |     genFn : (a -> Gen b) -> Gen (a :-> Cotree b)

438 |     genFn g = MkGen $ \sz, sd =>

439 |       iterate (shrinkFn forest) . map (runGen sz sd) $ build g

443 | apply : Fn a b -> a -> b

444 | apply (MkFn b f) = maybe b value . apply' f

454 | public export

455 | function' : ShrCogen a => Gen b -> Gen (a -> b)

456 | function' = map apply . function