
  0 | module Derive.Cogen
  1 | 
  2 | import public Hedgehog.Internal.Function as Hedgehog
  3 | 
  4 | import public Language.Reflection.Util
  5 | 
  6 | %default total
  7 | 
  8 | ||| Derivation facility for `Gogen` interface
  9 | |||
 10 | ||| Use `derive`, `deriveIndexed` or `derivePattern` from
 11 | ||| `Language.Reflection.Derive` for simple, purely indexed or mixed data types
 12 | export
 13 | CogenVis : Visibility -> List Name -> ParamTypeInfo -> Res (List TopLevel)
 14 | CogenVis vis nms p = do
 15 |   let fun  := funName p "perturb"
 16 |   let impl := implName p "Cogen"
 17 |   Right
 18 |     [ TL (perturbClaim fun p) (perturbDef fun p.info)
 19 |     , TL (cogenImplClaim impl p) (cogenImplDef fun impl)
 20 |     ]
 21 | 
 22 |   where
 23 |     perturbClaim : Name -> ParamTypeInfo -> Decl
 24 |     perturbClaim fun p =
 25 |       simpleClaim vis fun $
 26 |         piAll `(~(p.applied) -> StdGen -> StdGen) $ allImplicits p "Cogen"
 27 | 
 28 |     perturbDef : Name -> TypeInfo -> Decl
 29 |     perturbDef fun ti =
 30 |       def fun $ map clause $ [0 .. length ti.cons] `zip` ti.cons where
 31 | 
 32 |       clause : (Nat, Con _ _) -> Clause
 33 |       clause (idx, con) =
 34 |         accumArgs unerased (\x => `(~(var fun) ~x)) rhs arg con where
 35 | 
 36 |         arg : BoundArg 1 Unerased -> TTImp
 37 |         arg $ BA g [x] _ =
 38 |           assertIfRec nms g.type
 39 |             `(Function.perturb ~(var x) . Function.shiftArg)
 40 | 
 41 |         rhs : SnocList TTImp -> TTImp
 42 |         rhs = foldr (\l, r => `(~l . ~r))
 43 |                     `(System.Random.Pure.variant (fromInteger ~(primVal $ BI $ cast idx)))
 44 | 
 45 |     cogenImplClaim : Name -> ParamTypeInfo -> Decl
 46 |     cogenImplClaim impl p = implClaimVis vis impl $ implType "Cogen" p
 47 | 
 48 |     cogenImplDef : (fun, impl : Name) -> Decl
 49 |     cogenImplDef fun impl =
 50 |       def impl $ pure $ patClause (var impl) (var "MkCogen" `app` var fun)
 51 | 
 52 | 
 53 | ||| Alias for `CogenVis Public`
 54 | export %inline
 55 | Cogen : List Name -> ParamTypeInfo -> Res (List TopLevel)
 56 | Cogen = CogenVis Public
 57 |