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

Random generation of unique variables and unique containers.
-}

{-# LANGUAGE TupleSections #-}

module Clash.Hedgehog.Unique
  ( genUnique
  , genUniqMap
  , sampleUniqMap
  , sampleAnyUniqMap
  , Bias(..)
  , sampleUniqMapBiased
  ) where

import Control.Applicative (Alternative(empty))
import Data.Either (rights)
import Hedgehog (MonadGen, Range)
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range

import Clash.Core.HasType
import Clash.Core.Subst (aeqType)
import Clash.Core.Type
import Clash.Data.UniqMap (UniqMap)
import qualified Clash.Data.UniqMap as UniqMap
import Clash.Unique

import Clash.Hedgehog.Internal.Bias

genUnique :: forall m. MonadGen m => m Unique
genUnique :: m Unique
genUnique = Range Unique -> m Unique
forall (m :: Type -> Type). MonadGen m => Range Unique -> m Unique
Gen.int Range Unique
forall a. (Bounded a, Integral a) => Range a
Range.linearBounded

genUniqMap
  :: forall m k v
   . (MonadGen m, Uniquable k)
  => Range Int
  -> m k
  -> m v
  -> m (UniqMap v)
genUniqMap :: Range Unique -> m k -> m v -> m (UniqMap v)
genUniqMap Range Unique
range m k
genKey m v
genValue =
  [(k, v)] -> UniqMap v
forall a b. Uniquable a => [(a, b)] -> UniqMap b
UniqMap.fromList ([(k, v)] -> UniqMap v) -> m [(k, v)] -> m (UniqMap v)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Range Unique -> m (k, v) -> m [(k, v)]
forall (m :: Type -> Type) a.
MonadGen m =>
Range Unique -> m a -> m [a]
Gen.list Range Unique
range ((,) (k -> v -> (k, v)) -> m k -> m (v -> (k, v))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> m k
genKey m (v -> (k, v)) -> m v -> m (k, v)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> m v
genValue)

sampleAnyUniqMap
  :: forall m v
   . (Alternative m, MonadGen m, HasType v)
  => UniqMap v
  -> m (v, [Type])
sampleAnyUniqMap :: UniqMap v -> m (v, [Type])
sampleAnyUniqMap UniqMap v
xs =
  let xs' :: UniqMap v
xs' = (v -> Bool) -> UniqMap v -> UniqMap v
forall b. (b -> Bool) -> UniqMap b -> UniqMap b
UniqMap.filter (Bool -> Bool
not (Bool -> Bool) -> (v -> Bool) -> v -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Bool
isPolyTy (Type -> Bool) -> (v -> Type) -> v -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v -> Type
forall a. HasType a => a -> Type
coreTypeOf) UniqMap v
xs
   in if UniqMap v -> Bool
forall a. UniqMap a -> Bool
UniqMap.null UniqMap v
xs' then m (v, [Type])
forall (f :: Type -> Type) a. Alternative f => f a
empty else do
     v
x <- [v] -> m v
forall (f :: Type -> Type) (m :: Type -> Type) a.
(Foldable f, MonadGen m) =>
f a -> m a
Gen.element (UniqMap v -> [v]
forall b. UniqMap b -> [b]
UniqMap.elems UniqMap v
xs')
     let holes :: [Type]
holes = [Either TyVar Type] -> [Type]
forall a b. [Either a b] -> [b]
rights ([Either TyVar Type] -> [Type])
-> (([Either TyVar Type], Type) -> [Either TyVar Type])
-> ([Either TyVar Type], Type)
-> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Either TyVar Type], Type) -> [Either TyVar Type]
forall a b. (a, b) -> a
fst (([Either TyVar Type], Type) -> [Type])
-> ([Either TyVar Type], Type) -> [Type]
forall a b. (a -> b) -> a -> b
$ Type -> ([Either TyVar Type], Type)
splitFunForallTy (v -> Type
forall a. HasType a => a -> Type
coreTypeOf v
x)

     (v, [Type]) -> m (v, [Type])
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (v
x, [Type]
holes)

sampleUniqMap
  :: forall m v
   . (Alternative m, MonadGen m, HasType v)
  => (v -> Bool)
  -> Type
  -> UniqMap v
  -> m (v, [Type])
sampleUniqMap :: (v -> Bool) -> Type -> UniqMap v -> m (v, [Type])
sampleUniqMap v -> Bool
p Type
hole UniqMap v
xs =
  let xs' :: UniqMap (v, [Type])
xs' = (v -> Maybe (v, [Type])) -> UniqMap v -> UniqMap (v, [Type])
forall a b. (a -> Maybe b) -> UniqMap a -> UniqMap b
UniqMap.mapMaybe v -> Maybe (v, [Type])
findFit ((v -> Bool) -> UniqMap v -> UniqMap v
forall b. (b -> Bool) -> UniqMap b -> UniqMap b
UniqMap.filter v -> Bool
p UniqMap v
xs)
   in if UniqMap (v, [Type]) -> Bool
forall a. UniqMap a -> Bool
UniqMap.null UniqMap (v, [Type])
xs' then m (v, [Type])
forall (f :: Type -> Type) a. Alternative f => f a
empty else [(v, [Type])] -> m (v, [Type])
forall (f :: Type -> Type) (m :: Type -> Type) a.
(Foldable f, MonadGen m) =>
f a -> m a
Gen.element (UniqMap (v, [Type]) -> [(v, [Type])]
forall b. UniqMap b -> [b]
UniqMap.elems UniqMap (v, [Type])
xs')
 where
  findFit :: v -> Maybe (v, [Type])
findFit v
x =
    ([Type] -> (v, [Type])) -> Maybe [Type] -> Maybe (v, [Type])
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (v
x,) (Type -> Maybe [Type]
findFitArgs (v -> Type
forall a. HasType a => a -> Type
coreTypeOf v
x))

  -- NOTE [finding more complex fits]
  -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  -- This is not good enough. If I have a hole of type A -> B and I have
  -- a candidate of type forall a. a -> B, I will give up because it is not
  -- alpha equalivalent or a function type. I could return [Either TyVar Type]
  -- and include foralls, but this would still not fit polymorphic holes. For
  -- example, if I have the hole A -> B, and the candidate forall a. a -> B, it
  -- would:
  --
  --   1. not aeq to hole, add Left a to params
  --   2. not aeq to hole, add Right a to params
  --   3. not aeq to hole, discard
  --
  -- The correct approach to take here is to figure out which arguments need
  -- to be provided such that the hole and the type of the candidate can be
  -- unified. However, unification is (1) not provided by clash-lib currently
  -- and (2) very non-trivial to implement given we have -XTypeFamilies.
  findFitArgs :: Type -> Maybe [Type]
findFitArgs Type
a
    | Type -> Type -> Bool
aeqType Type
hole Type
a        = [Type] -> Maybe [Type]
forall a. a -> Maybe a
Just []
    | FunTy Type
b Type
c <- Type -> TypeView
tyView Type
a = ([Type] -> [Type]) -> Maybe [Type] -> Maybe [Type]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (Type
b Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:) (Type -> Maybe [Type]
findFitArgs Type
c)
    | Bool
otherwise             = Maybe [Type]
forall a. Maybe a
Nothing

sampleUniqMapBiased
  :: forall m v
   . (Alternative m, MonadGen m, HasType v, Bias v)
  => (v -> Bool)
  -> Type
  -> UniqMap v
  -> m (v, [Type])
sampleUniqMapBiased :: (v -> Bool) -> Type -> UniqMap v -> m (v, [Type])
sampleUniqMapBiased v -> Bool
p Type
hole UniqMap v
xs =
  let xs' :: [(v, [Type])]
xs' = UniqMap (v, [Type]) -> [(v, [Type])]
forall b. UniqMap b -> [b]
UniqMap.elems (UniqMap (v, [Type]) -> [(v, [Type])])
-> UniqMap (v, [Type]) -> [(v, [Type])]
forall a b. (a -> b) -> a -> b
$ (v -> Maybe (v, [Type])) -> UniqMap v -> UniqMap (v, [Type])
forall a b. (a -> Maybe b) -> UniqMap a -> UniqMap b
UniqMap.mapMaybe v -> Maybe (v, [Type])
findFit ((v -> Bool) -> UniqMap v -> UniqMap v
forall b. (b -> Bool) -> UniqMap b -> UniqMap b
UniqMap.filter v -> Bool
p UniqMap v
xs)
      bs :: [Unique]
bs  = ((v, [Type]) -> Unique) -> [(v, [Type])] -> [Unique]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (v -> Unique
forall a. Bias a => a -> Unique
biasOf (v -> Unique) -> ((v, [Type]) -> v) -> (v, [Type]) -> Unique
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (v, [Type]) -> v
forall a b. (a, b) -> a
fst) [(v, [Type])]
xs'
   in if [(v, [Type])] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
null [(v, [Type])]
xs' then m (v, [Type])
forall (f :: Type -> Type) a. Alternative f => f a
empty else [(Unique, m (v, [Type]))] -> m (v, [Type])
forall (m :: Type -> Type) a. MonadGen m => [(Unique, m a)] -> m a
Gen.frequency ([Unique] -> [m (v, [Type])] -> [(Unique, m (v, [Type]))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Unique]
bs ((v, [Type]) -> m (v, [Type])
forall (m :: Type -> Type) a. MonadGen m => a -> m a
Gen.constant ((v, [Type]) -> m (v, [Type])) -> [(v, [Type])] -> [m (v, [Type])]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> [(v, [Type])]
xs'))
  where
  findFit :: v -> Maybe (v, [Type])
findFit v
x =
    ([Type] -> (v, [Type])) -> Maybe [Type] -> Maybe (v, [Type])
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (v
x,) (Type -> Maybe [Type]
findFitArgs (v -> Type
forall a. HasType a => a -> Type
coreTypeOf v
x))

  findFitArgs :: Type -> Maybe [Type]
findFitArgs Type
a
    | Type -> Type -> Bool
aeqType Type
hole Type
a        = [Type] -> Maybe [Type]
forall a. a -> Maybe a
Just []
    | FunTy Type
b Type
c <- Type -> TypeView
tyView Type
a = ([Type] -> [Type]) -> Maybe [Type] -> Maybe [Type]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (Type
b Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:) (Type -> Maybe [Type]
findFitArgs Type
c)
    | Bool
otherwise             = Maybe [Type]
forall a. Maybe a
Nothing