{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE RankNTypes #-}

-- | 
-- Copyright: (c) Luca S. Jaekel
-- License: AGPL3
--
-- Infinitrees are memoization trees, which can be used to avoid dealing with mutable caches.

module Data.Infinitree
-- export the structure and the instances but not any accessors because you're not meant to invalidate the Infinitree
( Infinitree()
-- * Identity trees
, nats
, ints
, nums
-- * Construction Functions
, build
, buildInt
, buildNum
)
where

-- Distributive is a typeclass which allows you to use the law of distribution on functors
-- It is a superclass constraint for Representable, which is why I have to define it
import Data.Distributive (Distributive (distribute))

-- Representable functors allow indexing and construction from indices
-- they have a index type (Rep :: Type -> Type)
-- index takes an index and a structure, returns the element
-- tabulate calls your function with every index to build the functor
import Data.Functor.Rep (Representable, Rep, tabulate, index)

-- Natural numbers are [0..]
import Numeric.Natural (Natural)

-- ternary operator of sorts
import Data.Bool (bool)

-- | This tree is infinite, it doesn't end anywhere.
--
--  You can index into it infitely.
--
--  It always has a left and a right branch. Every Branch also holds a value.

data Infinitree a = Branch
  { left :: Infinitree a -- left branch, smaller number, the first left branch contains all odd numbers
  , leaf :: a -- current number, 0 for all intents and purposes
  , right :: Infinitree a -- right branch, bigger number, the first right branch contains all even numbers
  }

instance Functor Infinitree where
  fmap :: (a -> b) -> Infinitree a -> Infinitree b
  fmap f tree = Branch (fmap f (left tree)) (f $ leaf tree) (fmap f (right tree))

instance Applicative Infinitree where
  pure :: a -> Infinitree a
  pure e = Branch (pure e) e (pure e)
  (<*>) :: Infinitree (a -> b) -> Infinitree a -> Infinitree b
  (<*>) (Branch fl f fr) (Branch vl v vr) = Branch (fl <*> vl) (f v) (fr <*> vr)
  liftA2 :: (a -> b -> c) -> Infinitree a -> Infinitree b -> Infinitree c
  liftA2 f (Branch la va ra) (Branch lb vb rb) = Branch (liftA2 f la lb) (f va vb) (liftA2 f ra rb)
  (*>) :: Infinitree a -> Infinitree b -> Infinitree b
  (*>) = flip const
  (<*) :: Infinitree a -> Infinitree b -> Infinitree a
  (<*) = const

-- >>> [1, 2] <* [1, 2]
-- [1,1,2,2]

-- I could not define a useful Foldable instance
--
-- instance Foldable Infinitree where
--  foldMap :: Monoid m => (a -> m) -> Infinitree a -> m
--  foldMap f (Branch l v r) = f v

-- | This is a superclass constraint for representable, but it is entirely implementable from Representable
--
-- I now learned that I could have derived it via the Co newtype from Data.Functor.Rep
--
-- https://hackage-content.haskell.org/package/adjunctions-4.4.3/docs/Data-Functor-Rep.html#t:Co

instance Distributive Infinitree where
  distribute :: Functor f => f (Infinitree a) -> Infinitree (f a)
  distribute f = tabulate (\ i -> fmap (flip index i) f)

-- Representable allows indexing and construction
instance Representable Infinitree where
  -- only natural numbers index into this structure
  type Rep Infinitree = Natural

  tabulate :: (Rep Infinitree -> a) -> Infinitree a
  tabulate f' = let
    
    -- build a tree structure of numbers, like this
    --                _0
    --              _/  \_
    --            _/      \_
    --          _/          \_
    --         /              \
    --       _1                2_    
    --      /  \              /   \
    --     /    \            /     \
    --    /      \          /       \
    --   3        5        4         6
    --  / \      / \      / \       / \
    -- 7   11   9   13   8   12   10   14
    tabulate' :: (Rep Infinitree -> a) -> Rep Infinitree -> Natural -> Infinitree a
    tabulate' f !i !s = let -- keep the indices strict to avoid function application chains
      l = i + s
      r = l + s
      s' = 2 * s
      in Branch (tabulate' f l s') (f i) (tabulate' f r s')

    in tabulate' f' 0 1

  -- index into the tree structure recursively
  -- the current leaf always has value 0, the index will be adjusted along the way
  index :: Infinitree a -> Rep Infinitree -> a
  index t n = let

    -- inner recursive function
    index' !tree !0 = leaf tree
    index' !tree !i = index' subtree q
      where
        (!q, !r)  = pred i `quotRem` 2 -- q is strict to avoid useless function application delays
        !subtree = bool right left (r == 0) $ tree

    in index' t n

-- * Identity trees
-- 
-- These are probably not optimal for performance, since you always have two trees in memory if you `fmap` over them


-- | a tree of natural numbers.
--
-- a use case would be to `fmap` over it to transform it.
-- 
-- >>> map (index nats) [0..15]
-- [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]


nats :: Infinitree Natural
nats = tabulate @Infinitree id

-- | tree of integer numbers
--
-- in case you don't want to transform to integer for mapping

ints :: Infinitree Integer
ints = tabulate @Infinitree toInteger

-- | a tree of generic numbers
--
-- if you need a specific number type, make sure you don't use a bounded type, the tree is infinite

nums :: Num n => Infinitree n
nums = tabulate @Infinitree fromIntegral


-- * Construction functions

-- | build using the infinitree indices

build :: (Natural -> a) -> Infinitree a
build = tabulate @Infinitree

-- | build using arbitrary-width integers

buildInt :: (Integer -> a) -> Infinitree a
buildInt = tabulate @Infinitree . (. toInteger)

-- | build using whatever num type you need

buildNum :: Num n => (n -> a) -> Infinitree a
buildNum = tabulate @Infinitree . (. fromIntegral)