-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Biunfoldable
-- Copyright   :  (c) Sjoerd Visscher 2014
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  sjoerd@w3future.com
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Class of data structures with 2 type arguments that can be unfolded.
-----------------------------------------------------------------------------
{-# LANGUAGE Safe #-}
module Data.Biunfoldable
  (

  -- * Biunfoldable
    Biunfoldable(..)
  , biunfold_
  , biunfoldBF
  , biunfoldBF_

  -- ** Specific unfolds
  , biunfoldr
  , fromLists
  , randomDefault
  , arbitraryDefault

  )
  where

import Control.Applicative
import Data.Unfolder
import Data.Functor.Constant
import Control.Monad.Trans.State
import qualified System.Random as R
import Test.QuickCheck (Arbitrary(..), Gen, sized, resize)
import Data.Maybe

-- | Data structures with 2 type arguments (kind @* -> * -> *@) that can be unfolded.
--
-- For example, given a data type
--
-- > data Tree a b = Empty | Leaf a | Node (Tree a b) b (Tree a b)
--
-- a suitable instance would be
--
-- > instance Biunfoldable Tree where
-- >   biunfold fa fb = choose
-- >     [ pure Empty
-- >     , Leaf <$> fa
-- >     , Node <$> biunfold fa fb <*> fb <*> biunfold fa fb
-- >     ]
--
-- i.e. it follows closely the instance for 'Bitraversable', but instead of matching on an input value,
-- we 'choose' from a list of all cases.
class Biunfoldable t where
  -- | Given a way to generate elements, return a way to generate structures containing those elements.
  biunfold :: Unfolder f => f a -> f b -> f (t a b)

-- | Unfold the structure, always using @()@ as elements.
biunfold_ :: (Biunfoldable t, Unfolder f) => f (t () ())
biunfold_ :: f (t () ())
biunfold_ = f () -> f () -> f (t () ())
forall (t :: * -> * -> *) (f :: * -> *) a b.
(Biunfoldable t, Unfolder f) =>
f a -> f b -> f (t a b)
biunfold (() -> f ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) (() -> f ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

-- | Breadth-first unfold, which orders the result by the number of 'choose' calls.
biunfoldBF :: (Biunfoldable t, Unfolder f) => f a -> f b -> f (t a b)
biunfoldBF :: f a -> f b -> f (t a b)
biunfoldBF = (BFS f (t a b) -> f (t a b))
-> (BFS f a -> BFS f b -> BFS f (t a b)) -> f a -> f b -> f (t a b)
forall (t :: (* -> *) -> * -> *) (f :: * -> *) c a b.
(UnfolderTransformer t, Unfolder f) =>
(t f c -> f c) -> (t f a -> t f b -> t f c) -> f a -> f b -> f c
ala2 BFS f (t a b) -> f (t a b)
forall (f :: * -> *) x. Unfolder f => BFS f x -> f x
bfs BFS f a -> BFS f b -> BFS f (t a b)
forall (t :: * -> * -> *) (f :: * -> *) a b.
(Biunfoldable t, Unfolder f) =>
f a -> f b -> f (t a b)
biunfold

-- | Unfold the structure breadth-first, always using @()@ as elements.
biunfoldBF_ :: (Biunfoldable t, Unfolder f) => f (t () ())
biunfoldBF_ :: f (t () ())
biunfoldBF_ = BFS f (t () ()) -> f (t () ())
forall (f :: * -> *) x. Unfolder f => BFS f x -> f x
bfs BFS f (t () ())
forall (t :: * -> * -> *) (f :: * -> *).
(Biunfoldable t, Unfolder f) =>
f (t () ())
biunfold_

-- | @biunfoldr@ builds a data structure from a seed value.
biunfoldr :: Biunfoldable t => (c -> Maybe (a, c)) -> (c -> Maybe (b, c)) -> c -> Maybe (t a b)
biunfoldr :: (c -> Maybe (a, c)) -> (c -> Maybe (b, c)) -> c -> Maybe (t a b)
biunfoldr c -> Maybe (a, c)
fa c -> Maybe (b, c)
fb c
z = [(t a b, c)] -> Maybe (t a b)
forall a. [(a, c)] -> Maybe a
terminate ([(t a b, c)] -> Maybe (t a b))
-> (StateT c [] (t a b) -> [(t a b, c)])
-> StateT c [] (t a b)
-> Maybe (t a b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StateT c [] (t a b) -> c -> [(t a b, c)])
-> c -> StateT c [] (t a b) -> [(t a b, c)]
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT c [] (t a b) -> c -> [(t a b, c)]
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT c
z (StateT c [] (t a b) -> Maybe (t a b))
-> StateT c [] (t a b) -> Maybe (t a b)
forall a b. (a -> b) -> a -> b
$ StateT c [] a -> StateT c [] b -> StateT c [] (t a b)
forall (t :: * -> * -> *) (f :: * -> *) a b.
(Biunfoldable t, Unfolder f) =>
f a -> f b -> f (t a b)
biunfoldBF ((c -> [(a, c)]) -> StateT c [] a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((c -> [(a, c)]) -> StateT c [] a)
-> (c -> [(a, c)]) -> StateT c [] a
forall a b. (a -> b) -> a -> b
$ Maybe (a, c) -> [(a, c)]
forall a. Maybe a -> [a]
maybeToList (Maybe (a, c) -> [(a, c)]) -> (c -> Maybe (a, c)) -> c -> [(a, c)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. c -> Maybe (a, c)
fa) ((c -> [(b, c)]) -> StateT c [] b
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((c -> [(b, c)]) -> StateT c [] b)
-> (c -> [(b, c)]) -> StateT c [] b
forall a b. (a -> b) -> a -> b
$ Maybe (b, c) -> [(b, c)]
forall a. Maybe a -> [a]
maybeToList (Maybe (b, c) -> [(b, c)]) -> (c -> Maybe (b, c)) -> c -> [(b, c)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. c -> Maybe (b, c)
fb)
  where
    terminate :: [(a, c)] -> Maybe a
terminate [] = Maybe a
forall a. Maybe a
Nothing
    terminate ((a
t, c
c):[(a, c)]
ts) = if Maybe (a, c) -> Bool
forall a. Maybe a -> Bool
isNothing (c -> Maybe (a, c)
fa c
c) Bool -> Bool -> Bool
&& Maybe (b, c) -> Bool
forall a. Maybe a -> Bool
isNothing (c -> Maybe (b, c)
fb c
c) then a -> Maybe a
forall a. a -> Maybe a
Just a
t else [(a, c)] -> Maybe a
terminate [(a, c)]
ts

-- | Create a data structure using the lists as input.
-- This can fail because there might not be a data structure with the same number
-- of element positions as the number of elements in the lists.
fromLists :: Biunfoldable t => [a] -> [b] -> Maybe (t a b)
fromLists :: [a] -> [b] -> Maybe (t a b)
fromLists = (([a], [b]) -> Maybe (t a b)) -> [a] -> [b] -> Maybe (t a b)
forall a b c. ((a, b) -> c) -> a -> b -> c
curry ((([a], [b]) -> Maybe (t a b)) -> [a] -> [b] -> Maybe (t a b))
-> (([a], [b]) -> Maybe (t a b)) -> [a] -> [b] -> Maybe (t a b)
forall a b. (a -> b) -> a -> b
$ (([a], [b]) -> Maybe (a, ([a], [b])))
-> (([a], [b]) -> Maybe (b, ([a], [b])))
-> ([a], [b])
-> Maybe (t a b)
forall (t :: * -> * -> *) c a b.
Biunfoldable t =>
(c -> Maybe (a, c)) -> (c -> Maybe (b, c)) -> c -> Maybe (t a b)
biunfoldr ([a], [b]) -> Maybe (a, ([a], [b]))
forall a b. ([a], b) -> Maybe (a, ([a], b))
unconsA ([a], [b]) -> Maybe (b, ([a], [b]))
forall a a. (a, [a]) -> Maybe (a, (a, [a]))
unconsB
  where
    unconsA :: ([a], b) -> Maybe (a, ([a], b))
unconsA ([], b
_) = Maybe (a, ([a], b))
forall a. Maybe a
Nothing
    unconsA (a
a:[a]
as, b
bs) = (a, ([a], b)) -> Maybe (a, ([a], b))
forall a. a -> Maybe a
Just (a
a, ([a]
as, b
bs))
    unconsB :: (a, [a]) -> Maybe (a, (a, [a]))
unconsB (a
_, []) = Maybe (a, (a, [a]))
forall a. Maybe a
Nothing
    unconsB (a
as, a
b:[a]
bs) = (a, (a, [a])) -> Maybe (a, (a, [a]))
forall a. a -> Maybe a
Just (a
b, (a
as, [a]
bs))

-- | Generate a random value, can be used as default instance for 'R.Random'.
randomDefault :: (R.Random a, R.Random b, R.RandomGen g, Biunfoldable t) => g -> (t a b, g)
randomDefault :: g -> (t a b, g)
randomDefault = State g (t a b) -> g -> (t a b, g)
forall s a. State s a -> s -> (a, s)
runState (State g (t a b) -> g -> (t a b, g))
-> (Random g Identity (t a b) -> State g (t a b))
-> Random g Identity (t a b)
-> g
-> (t a b, g)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Random g Identity (t a b) -> State g (t a b)
forall g (m :: * -> *) a. Random g m a -> StateT g m a
getRandom (Random g Identity (t a b) -> g -> (t a b, g))
-> Random g Identity (t a b) -> g -> (t a b, g)
forall a b. (a -> b) -> a -> b
$ Random g Identity a
-> Random g Identity b -> Random g Identity (t a b)
forall (t :: * -> * -> *) (f :: * -> *) a b.
(Biunfoldable t, Unfolder f) =>
f a -> f b -> f (t a b)
biunfold (StateT g Identity a -> Random g Identity a
forall g (m :: * -> *) a. StateT g m a -> Random g m a
Random (StateT g Identity a -> Random g Identity a)
-> ((g -> (a, g)) -> StateT g Identity a)
-> (g -> (a, g))
-> Random g Identity a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (g -> (a, g)) -> StateT g Identity a
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((g -> (a, g)) -> Random g Identity a)
-> (g -> (a, g)) -> Random g Identity a
forall a b. (a -> b) -> a -> b
$ g -> (a, g)
forall a g. (Random a, RandomGen g) => g -> (a, g)
R.random) (StateT g Identity b -> Random g Identity b
forall g (m :: * -> *) a. StateT g m a -> Random g m a
Random (StateT g Identity b -> Random g Identity b)
-> ((g -> (b, g)) -> StateT g Identity b)
-> (g -> (b, g))
-> Random g Identity b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (g -> (b, g)) -> StateT g Identity b
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((g -> (b, g)) -> Random g Identity b)
-> (g -> (b, g)) -> Random g Identity b
forall a b. (a -> b) -> a -> b
$ g -> (b, g)
forall a g. (Random a, RandomGen g) => g -> (a, g)
R.random)

-- | Provides a QuickCheck generator, can be used as default instance for 'Arbitrary'.
arbitraryDefault :: (Arbitrary a, Arbitrary b, Biunfoldable t) => Gen (t a b)
arbitraryDefault :: Gen (t a b)
arbitraryDefault = let Arb Int
_ Int
_ Gen (Maybe (t a b))
gen = Arb a -> Arb b -> Arb (t a b)
forall (t :: * -> * -> *) (f :: * -> *) a b.
(Biunfoldable t, Unfolder f) =>
f a -> f b -> f (t a b)
biunfold Arb a
forall a. Arbitrary a => Arb a
arbUnit Arb b
forall a. Arbitrary a => Arb a
arbUnit in
  t a b -> Maybe (t a b) -> t a b
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> t a b
forall a. HasCallStack => [Char] -> a
error [Char]
"Failed to generate a value.") (Maybe (t a b) -> t a b) -> Gen (Maybe (t a b)) -> Gen (t a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Maybe (t a b))
gen

instance Biunfoldable Either where
  biunfold :: f a -> f b -> f (Either a b)
biunfold f a
fa f b
fb = [f (Either a b)] -> f (Either a b)
forall (f :: * -> *) a. Unfolder f => [f a] -> f a
choose
    [ a -> Either a b
forall a b. a -> Either a b
Left (a -> Either a b) -> f a -> f (Either a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
fa
    , b -> Either a b
forall a b. b -> Either a b
Right (b -> Either a b) -> f b -> f (Either a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f b
fb
    ]

instance Biunfoldable (,) where
  biunfold :: f a -> f b -> f (a, b)
biunfold f a
fa f b
fb = [f (a, b)] -> f (a, b)
forall (f :: * -> *) a. Unfolder f => [f a] -> f a
choose
    [ (,) (a -> b -> (a, b)) -> f a -> f (b -> (a, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
fa f (b -> (a, b)) -> f b -> f (a, b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> f b
fb ]

instance Biunfoldable Constant where
  biunfold :: f a -> f b -> f (Constant a b)
biunfold f a
fa f b
_ = [f (Constant a b)] -> f (Constant a b)
forall (f :: * -> *) a. Unfolder f => [f a] -> f a
choose
    [ a -> Constant a b
forall k a (b :: k). a -> Constant a b
Constant (a -> Constant a b) -> f a -> f (Constant a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
fa ]