-- CANNOT TCO

-- The recursive call to foldr' fn acc xs is inside an argument to `fn`.
-- After the inner call to foldr' returns, you still need to compute fn.
-- There is work left to do after the recursive call, thus the call isn't in
-- tail position.
foldr' :: ∀ a b. (a -> b -> b) -> b -> List a -> b
foldr' _ acc Nil = acc
foldr' fn acc (Cons x xs) = fn x (foldr' fn acc xs)

-- TCO Possible
-- The recursive call is the final action. You compute the accumulator value first
-- and then make a tail call to foldl'. No extra work remains
foldl' :: ∀ a b. (b -> a -> b) -> b -> List a -> b
foldl' _ acc Nil = acc
foldl' fn acc (Cons x xs) = foldl' fn (fn acc x) xs

foldrCPS :: ∀ a b. (a -> b -> b) -> b -> List a -> b
foldrCPS fn acc lst = go lst identity
  where
  go Nil k = k acc
  go (Cons x xs) k = go xs (\res -> k (fn x res))

> foldrCPS (\x xs -> x + xs) 0 $ List.range 0 100000
file:///Users/nathanwaltz/Documents/Prototype/.psci_modules/$PSCI/index.js:5
var it = /* #__PURE__ */ Rose["foldr$prime$prime"](function (x) {
                                                            ^

RangeError: Maximum call stack size exceeded

tailRec :: forall a b. (a -> Step a b) -> a -> b

-- example
pow :: Int -> Int -> Int
pow n p = tailRec go { accum: 1, power: p }
  where
  go :: _ -> Step _ Int
  go { accum: acc, power: 0 } = Done acc
  go { accum: acc, power: p } = Loop { accum: acc * n, power: p - 1 }

data Step a b = Loop a | Done b

newtype State a b = State
  { k :: b -> Step (State a b ) b
  , lst :: List a
  }

foldrCPS' :: ∀ a b. (a -> b -> b) -> b -> List a -> b
foldrCPS' fn acc lst = tailRec go (State { lst, k: Done })
  where
  go (State { lst: Nil, k }) = k acc
  go (State { lst: Cons x xs, k }) = Loop (State { lst: xs, k: \res -> k (fn x res) })

-- making the continuation explicit
foldrCPS'' :: ∀ a b. (a -> b -> b) -> b -> List a -> b
foldrCPS'' fn acc lst = go lst identity
  where
  go Nil k = k acc
  go (Cons x xs) k = go xs (goCont x k)

  goCont currentValue kont = \res ->
    kont (fn currentValue res)

-- encode both the goCont and the identity function as data
data FoldAccum a b
  = FoldCont a (FoldAccum a b)
  | FAContIdentity

foldrCPS :: ∀ a b. (a -> b -> b) -> b -> List a -> b
foldrCPS fn acc lst = go lst FAContIdentity
  where
  go Nil k = eval k acc
  go (Cons x xs) k = go xs (FoldCont x k)

  -- small step evaluation
  eval cont acc' = case cont of
    FoldCont a next -> do
      let result = fn a acc'
      eval next result
    FAContIdentity ->
      acc'

-- naive implementation
fib :: Int -> Maybe Int
fib 0 = Just 0
fib 1 = Just 1
fib n = fib (n - 1) + fib (n - 2)

-- CPS fib
fib :: Int -> Maybe Int
fib n = fib' identity
  where
  fib' k 0 = k $ Just 0
  fib' k 1 = k $ Just 1
  fib' k n
   | n < 0 = k Nothing
   | otherwise =
      fib'
        (\a ->
          fib'
            (\b ->
              k $ (+) <$> a <*> b
            ) (n - 1)
        ) (n - 2)

-- CPS fib - explicit continuations
fib :: Int -> Maybe Int
fib = fib' identity
  where
  fib' k 0 = k $ Just 0
  fib' k 1 = k $ Just 1
  fib' k n
   | n < 0 = k Nothing
   | otherwise = fib' (contFibOne n k) (n - 2)

  contFibOne n cont = \a -> do
    fib' (contFibTwo a cont) (n - 1)

  contFibTwo a cont = \b ->
    cont ((+) <$> a <*> b)

data FibKont
  = ContFibOne Int FibKont
  | ContFibTwo (Maybe Int) FibKont
  | FibIdentity

-- CPS fib - defunctionalization
fib :: Int -> Maybe Int
fib = fib' FibIdentity
  where
  fib' k 0 = eval k (Just 0)
  fib' k 1 = eval k (Just 1)
  fib' k n
   | n < 0 = eval k Nothing
   | otherwise = fib' (ContFibOne n k) (n - 2)

  eval :: FibKont -> Maybe Int -> Maybe Int
  eval cont acc = case cont of
    ContFibOne n next ->
      fib' (ContFibTwo acc next) (n - 1)
    ContFibTwo x next ->
      eval next ((+) <$> x <*> acc)
    FibIdentity -> acc

-- now we need to turn mutually recursive fib', eval calls
-- into a data type. This is great because we can now do small step
-- evaluation, and go is always in the tail position. We are forcing heap
-- allocations with the explicit FibCall and FibKont ADTs. These heap
-- allocations are what we are trading stack allocations for,
-- and we get small step evaluation of a tree-shaped recursive function.
-- Recursion thus becomes iteration!

data FibKont
  = ContFibOne Int FibKont
  | ContFibTwo (Maybe Int) FibKont
  | FibIdentity

fib :: Int -> Maybe Int
fib x = go (FibFib' FibIdentity x)
  where
  go call = case call of
    FibFib' cont 0 -> go (FibEval cont (Just 0))
    FibFib' cont 1 -> go (FibEval cont (Just 1))
    FibFib' cont n | n < 0 -> go (FibEval cont Nothing)
    FibFib' cont n -> go (FibFib' (ContFibOne n cont) (n - 2))
    FibEval cont acc -> case cont of
        ContFibOne n next ->
          go (FibFib' (ContFibTwo acc next) (n - 1))
        ContFibTwo a next ->
          go (FibEval next ((+) <$> a <*> acc))
        FibIdentity -> acc

data FibCall
  = FibEval FibKont (Maybe Int)
  | FibFib' FibKont Int

-- Mu over the composition of thunks
type DelayedMu f = Mu (Compose (Function Unit) f)

-- Run a thunked small-step fixed point.
fix :: ∀ b. Mu (Compose (Function Unit) (Either b)) -> b
fix mu = case unroll mu of
  Compose thunk ->
    case thunk unit of
      Left result -> result
      Right mu' -> fix mu'

fib :: Int -> Maybe Int
fib x = fix $ go (FibFib' FibIdentity x)
  where
  go call = roll $ Compose $ \_ -> case call of
    FibFib' cont 0 -> Right $ go (FibEval cont (Just 0))
    FibFib' cont 1 -> Right $ go (FibEval cont (Just 1))
    FibFib' cont n | n < 0 -> Right $ go (FibEval cont Nothing)
    FibFib' cont n -> Right $ go (FibFib' (ContFibOne n cont) (n - 2))
    FibEval cont acc -> case cont of
        ContFibOne n next ->
          Right $ go (FibFib' (ContFibTwo acc next) (n - 1))
        ContFibTwo a next ->
          Right $ go (FibEval next ((+) <$> a <*> acc))
        FibIdentity -> Left acc

data Free f b
  = Pure b
  | Free (f (Free f b))

type Trampoline = Free ((->) Unit)

-- we create a bunch of thunks of free, and then lift those into free
suspend :: ∀ b. (Unit -> Free ((->) Unit) b) → Free ((->) Unit) b
suspend = join <<< delay

-- probably the most concise - notice the do notation!
fibTrampoline :: Int -> Maybe Int
fibTrampoline n = runTrampoline (fib' n)
  where
  fib' 0 = suspend \_ -> pure $ Just 0
  fib' 1 = suspend \_ -> pure $ Just 1
  fib' x
   | x < 0 = suspend \_ -> pure Nothing
   | otherwise = suspend \_ -> do
      a' <- fib' (x - 2)
      b' <- fib' (x - 1)
      pure ((+) <$> a' <*> b')