selective-countable.hs

-- Selective lets us do "monadic" binds where the domain is restricted
-- to the countable types (which can be infinite).

{-# LANGUAGE TypeOperators     #-}
{-# LANGUAGE PatternSynonyms   #-}
{-# LANGUAGE ViewPatterns      #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts  #-}

import GHC.Generics
import Data.Coerce
import Control.Selective
import Control.Applicative
import Numeric.Natural

----------------------------------------------------------------------
-- Countable class
----------------------------------------------------------------------
  
class Countable' f where enumerated' :: [f a]
  
instance Countable' V1 where enumerated' = []
instance Countable' U1 where enumerated' = [U1]

interleaving :: (a -> c) -> (b -> c) -> [a] -> [b] -> [c]  
interleaving f g (x:xs) ys = f x : interleaving g f ys xs
interleaving f g []     ys = map g ys
  
instance (Countable' f, Countable' g) => Countable' (f :+: g) where
  enumerated' = interleaving L1 R1 enumerated' enumerated'

convolving :: (a -> b -> c) -> [a] -> [b] -> [c]
convolving c xs ys = concat (foldr f [] xs) where
  f x = foldr g id ys . ([] :) where
    g y k (uncons -> ~(z, zs)) = (c x y : z) : k zs
  uncons []     = ([],[])
  uncons (x:xs) = (x,xs)
  
instance (Countable' f, Countable' g) => Countable' (f :*: g) where
  enumerated' = convolving (:*:) enumerated' enumerated'

instance Countable c => Countable' (K1 i c) where
  enumerated' = (coerce :: [c] -> [K1 i c a]) enumerated
  
instance Countable' f => Countable' (M1 i t f) where
  enumerated' = (coerce :: [f a] -> [M1 i t f a]) enumerated'
  
class Countable a where
  enumerated :: [a]
  default enumerated :: (Generic a, Countable' (Rep a)) => [a]
  enumerated = map to enumerated'

instance Countable ()
instance Countable Bool
instance Countable a => Countable [a]
instance Countable Natural where enumerated = [0..]
instance Countable Integer where enumerated = interleaving id negate [0..] [1..]

----------------------------------------------------------------------
-- Non-monadic binds
----------------------------------------------------------------------

(>>#=) :: (Eq a, Countable a, Selective f) => f a -> (a -> f b) -> f b
xs >>#= k = foldr f undefined enumerated
  where
    f y = ifS (fmap (==y) xs) (k y)
    
(>>!=) :: (Eq a, Countable a, Applicative f) => f a -> (a -> f b) -> f b
xs >>!= k = foldr f undefined enumerated
  where
    f y = liftA3 (\b x y -> if b then x else y) (fmap (==y) xs) (k y)

----------------------------------------------------------------------
-- Examples
----------------------------------------------------------------------

-- Selective lets us do a monadic bind without running the
-- effects on every branch, even when the domain is infinite (but
-- countably so).

-- Normal bind:
--
-- >>> pure 1 >>= print
-- 1

-- Selective:
--
-- >>> pure 1 >>#= print
-- 1

-- Applicative only:
--
-- >>> pure 1 >>!= print
-- 0
-- -1
-- 1
-- -2
-- 2
-- ...

-- This also means we can do binds over infinite (but countable)
-- domains even when we don't care about "effects".

-- Normal bind:
--
-- >>> Just 1 >>= Just
-- Just 1

-- Selective:
--
-- >>> Just 1 >>#= Just
-- Just 1

-- Applicative only:
--
-- >>> Just 1 >>!= Just
-- _|_