{-|
Copyright   : (C) 2021, QBayLogic B.V.
License     : BSD2 (see the file LICENSE)
Maintainer  : QBayLogic B.V. <devops@qbaylogic.com>

Random generation of core variables.
-}

module Clash.Hedgehog.Core.Var
  ( genTyVar
  , genId
  , genLocalId
  , genGlobalId
  , genVars
  ) where

import Hedgehog (MonadGen)
import qualified Hedgehog.Gen as Gen

import Clash.Core.Name (Name(nameUniq))
import Clash.Core.Term (TmName)
import Clash.Core.Type (Kind, KindOrType, TyName, Type)
import Clash.Core.Var (Id, IdScope(..), TyVar, Var(..))
import qualified Clash.Data.UniqMap as UniqMap

import Clash.Hedgehog.Core.Name (genFreshName)

-- | Generate a fresh type variable of the specified kind.
genTyVar :: forall m. MonadGen m => Kind -> m TyName -> m TyVar
genTyVar :: Kind -> m TyName -> m TyVar
genTyVar Kind
kn m TyName
genName = do
  TyName
name <- m TyName
genName
  TyVar -> m TyVar
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (TyName -> Unique -> Kind -> TyVar
forall a. Name a -> Unique -> Kind -> Var a
TyVar TyName
name (TyName -> Unique
forall a. Name a -> Unique
nameUniq TyName
name) Kind
kn)

-- | Generate a fresh identifier of the specified kind.
genId :: forall m. MonadGen m => Type -> m TmName -> m Id
genId :: Kind -> m TmName -> m Id
genId Kind
ty m TmName
genName = do
  TmName
name  <- m TmName
genName
  IdScope
scope <- [IdScope] -> m IdScope
forall (m :: Type -> Type) a. MonadGen m => [a] -> m a
Gen.element [IdScope
GlobalId, IdScope
LocalId]
  Id -> m Id
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (TmName -> Unique -> Kind -> IdScope -> Id
forall a. Name a -> Unique -> Kind -> IdScope -> Var a
Id TmName
name (TmName -> Unique
forall a. Name a -> Unique
nameUniq TmName
name) Kind
ty IdScope
scope)

-- | Generate a fresh local identifier of the specified kind.
genLocalId :: forall m. MonadGen m => Type -> m TmName -> m Id
genLocalId :: Kind -> m TmName -> m Id
genLocalId Kind
ty =
  (Id -> Id) -> m Id -> m Id
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Id -> Id
forall a. Var a -> Var a
setToLocal (m Id -> m Id) -> (m TmName -> m Id) -> m TmName -> m Id
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Kind -> m TmName -> m Id
forall (m :: Type -> Type). MonadGen m => Kind -> m TmName -> m Id
genId Kind
ty
 where
  setToLocal :: Var a -> Var a
setToLocal i :: Var a
i@Id{} = Var a
i {idScope :: IdScope
idScope = IdScope
LocalId}
  setToLocal Var a
i = Var a
i

-- | Generate a fresh global identifier of the specified kind.
genGlobalId :: forall m. MonadGen m => Type -> m TmName -> m Id
genGlobalId :: Kind -> m TmName -> m Id
genGlobalId Kind
ty =
  (Id -> Id) -> m Id -> m Id
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Id -> Id
forall a. Var a -> Var a
setToGlobal (m Id -> m Id) -> (m TmName -> m Id) -> m TmName -> m Id
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Kind -> m TmName -> m Id
forall (m :: Type -> Type). MonadGen m => Kind -> m TmName -> m Id
genId Kind
ty
 where
  setToGlobal :: Var a -> Var a
setToGlobal i :: Var a
i@Id{} = Var a
i {idScope :: IdScope
idScope = IdScope
LocalId}
  setToGlobal Var a
i = Var a
i

mapAccumLM
  :: forall m acc x y
   . Monad m
  => (acc -> x -> m (acc, y))
  -> acc
  -> [x]
  -> m (acc, [y])
mapAccumLM :: (acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM acc -> x -> m (acc, y)
_ acc
acc [] = (acc, [y]) -> m (acc, [y])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (acc
acc, [])
mapAccumLM acc -> x -> m (acc, y)
f acc
acc (x
x:[x]
xs) = do
  (acc
acc', y
y) <- acc -> x -> m (acc, y)
f acc
acc x
x
  (acc
acc'', [y]
ys) <- (acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
forall (m :: Type -> Type) acc x y.
Monad m =>
(acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM acc -> x -> m (acc, y)
f acc
acc' [x]
xs
  (acc, [y]) -> m (acc, [y])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (acc
acc'', y
yy -> [y] -> [y]
forall a. a -> [a] -> [a]
:[y]
ys)

-- | Generate a collection of variables, from a supplied function to generate
-- variables and the kinds / types of variables to generate.
--
-- TODO While this gives "unique" vars because the uniques are different, it
-- can generate multiple vars with the same OccName.
genVars
  :: forall m a
   . MonadGen m
  => (KindOrType -> m (Name a) -> m (Var a))
  -> [KindOrType]
  -> m (Name a)
  -> m [Var a]
genVars :: (Kind -> m (Name a) -> m (Var a))
-> [Kind] -> m (Name a) -> m [Var a]
genVars Kind -> m (Name a) -> m (Var a)
genVar [Kind]
kts m (Name a)
genName =
  (UniqMap (Var a), [Var a]) -> [Var a]
forall a b. (a, b) -> b
snd ((UniqMap (Var a), [Var a]) -> [Var a])
-> m (UniqMap (Var a), [Var a]) -> m [Var a]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (UniqMap (Var a) -> Kind -> m (UniqMap (Var a), Var a))
-> UniqMap (Var a) -> [Kind] -> m (UniqMap (Var a), [Var a])
forall (m :: Type -> Type) acc x y.
Monad m =>
(acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM UniqMap (Var a) -> Kind -> m (UniqMap (Var a), Var a)
go UniqMap (Var a)
forall a. Monoid a => a
mempty [Kind]
kts
 where
  go :: UniqMap (Var a) -> Kind -> m (UniqMap (Var a), Var a)
go UniqMap (Var a)
used Kind
kt = do
    Var a
var <- Kind -> m (Name a) -> m (Var a)
genVar Kind
kt (UniqMap (Var a) -> m (Name a) -> m (Name a)
forall (m :: Type -> Type) a b.
MonadGen m =>
UniqMap b -> m (Name a) -> m (Name a)
genFreshName UniqMap (Var a)
used m (Name a)
genName)
    (UniqMap (Var a), Var a) -> m (UniqMap (Var a), Var a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Var a -> UniqMap (Var a) -> UniqMap (Var a)
forall a. Uniquable a => a -> UniqMap a -> UniqMap a
UniqMap.insertUnique Var a
var UniqMap (Var a)
used, Var a
var)