Iteration.hs

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}

-- Inspired by:
-- https://github.com/fumieval/data-functor-logistic
-- https://gist.github.com/viercc/2e6c1d8566a6fbaf0d21c09103e60b76

module Iteration where

import Data.Distributive
import Numeric.Natural
import Prelude hiding (iterate)
import Control.Applicative (Const(..))
import Data.Traversable (mapAccumL)
import qualified Data.List as List (iterate)

getters :: Distributive t => t (t a -> a)
getters = distribute id

class Iteration t where
    iterate :: (a -> a) -> a -> t a

-- A plausible law: given Foldable t,
-- toList (iterate f i) = take (length getters) (List.iterate f i)

nats :: Iteration t => t (Natural)
nats = iterate succ 0

-- cf. Data.Functor.Rep.distributeRepBy
gridBy :: Distributive t => (a -> a -> b) -> t a -> t (t b)
gridBy f u = (\g -> (\h -> g u `f` h u) <$> getters) <$> getters

diag :: (Distributive t, Iteration t) => t (t Bool)
diag = gridBy (==) nats

-- Extra stuff:

pureD :: Distributive t => a -> t a
pureD = cotraverse getConst . Const

iterateDT :: (Distributive t, Traversable t) => (a -> a) -> a -> t a
iterateDT f i = (snd . mapAccumL (\a _ -> (f a, a)) i) (pureD ())

-- Like a stream unfold, though possibly not the anamorphism for t a.
unfoldI :: (Functor t, Iteration t) => (b -> (a, b)) -> b -> t a
unfoldI f = fmap fst . iterate (f . snd) . f

-- Example instances:

instance Iteration [] where
    iterate = List.iterate

data Duo a = Duo { fstDuo :: a, sndDuo :: a }
    deriving (Show, Functor, Foldable, Traversable)

instance Distributive Duo where
    distribute u = Duo (fstDuo <$> u) (sndDuo <$> u)

instance Iteration Duo where
    iterate f i = Duo i (f i)