import Control.Applicative (empty, liftA2)

import Control.Monad (guard)

import Control.Monad.State (State, evalState, get, modify)

import Data.Foldable (foldr)

import Data.Functor (($>), (<&>))

import qualified Data.Set as S

import Data.Text (Text, pack)

tco = flip evalState 0 . everywhereTopDownM convert where

tcoDoneM :: State Int Text

tcoDoneM = get <&> \count -> "$tco_done" <>

if count == 0 then "" else pack . show $ count

convert :: AST -> State Int AST

| Just trFns <- findTailRecursiveFns name arity body'

= VariableIntroduction ss name . Just . replace <$> toLoop trFns name arity outerArgs innerArgs body'

arity = length argss

-- ^ this is the number of calls, not the number of arguments, if there's

-- ever a practical difference.

(argss, body', replace) = topCollectAllFunctionArgs [] id fn

convert js = pure js

rewriteFunctionsWith :: ([Text] -> [Text]) -> [[Text]] -> (AST -> AST) -> AST -> ([[Text]], AST, AST -> AST)

rewriteFunctionsWith argMapper = collectAllFunctionArgs

where

collectAllFunctionArgs allArgs f (Function s1 ident args (Block s2 (body@(Return _ _):_))) =

collectAllFunctionArgs (args : allArgs) (\b -> f (Function s1 ident (argMapper args) (Block s2 [b]))) body

collectAllFunctionArgs allArgs f (Function ss ident args body@(Block _ _)) =

(args : allArgs, body, f . Function ss ident (argMapper args))

collectAllFunctionArgs allArgs f (Return s1 (Function s2 ident args (Block s3 [body]))) =

collectAllFunctionArgs (args : allArgs) (\b -> f (Return s1 (Function s2 ident (argMapper args) (Block s3 [b])))) body

collectAllFunctionArgs allArgs f (Return s1 (Function s2 ident args body@(Block _ _))) =

(args : allArgs, body, f . Return s1 . Function s2 ident (argMapper args))

collectAllFunctionArgs allArgs f body = (allArgs, body, f)

topCollectAllFunctionArgs :: [[Text]] -> (AST -> AST) -> AST -> ([[Text]], AST, AST -> AST)

topCollectAllFunctionArgs = rewriteFunctionsWith (map copyVar)

innerCollectAllFunctionArgs :: [[Text]] -> (AST -> AST) -> AST -> ([[Text]], AST, AST -> AST)

innerCollectAllFunctionArgs = rewriteFunctionsWith id

countReferences :: Text -> AST -> Int

countReferences ident = everything (+) match where

match :: AST -> Int

match (Var _ ident') | ident == ident' = 1

match _ = 0

-- If `ident` is a tail-recursive function, returns a set of identifiers

-- that are locally bound to functions participating in the tail recursion.

-- Otherwise, returns Nothing.

findTailRecursiveFns :: Text -> Int -> AST -> Maybe (S.Set Text)

findTailRecursiveFns ident arity js = guard (countReferences ident js > 0) *> go (S.empty, S.singleton (ident, arity))

where

go :: (S.Set Text, S.Set (Text, Int)) -> Maybe (S.Set Text)

go (known, required) =

case S.minView required of

Just (r, required') -> do

required'' <- findTailPositionDeps r js

go (S.insert (fst r) known, required' <> (S.filter (not . (`S.member` known) . fst) required''))

Nothing ->

pure known

-- Returns set of identifiers (with their arities) that need to be used

-- exclusively in tail calls using their full arity in order for this

-- identifier to be considered in tail position (or Nothing if this

-- identifier is used somewhere not as a tail call with full arity).

findTailPositionDeps :: (Text, Int) -> AST -> Maybe (S.Set (Text, Int))

findTailPositionDeps (ident, arity) js = allInTailPosition js where

countSelfReferences = countReferences ident

| isSelfCall ident arity expr = guard (countSelfReferences expr == 1) $> S.empty

| otherwise = guard (countSelfReferences expr == 0) $> S.empty

= guard (countSelfReferences js1 == 0) *> allInTailPosition body

= guard (countSelfReferences js1 == 0 && countSelfReferences js2 == 0) *> allInTailPosition body

= guard (countSelfReferences js1 == 0) *> allInTailPosition body

= guard (countSelfReferences js1 == 0) *> liftA2 mappend (allInTailPosition body) (foldMapA allInTailPosition el)

= foldMapA allInTailPosition body

= guard (countSelfReferences js1 == 0) $> S.empty

= pure S.empty

allInTailPosition (VariableIntroduction _ _ Nothing)

= pure S.empty

allInTailPosition (VariableIntroduction _ ident' (Just js1))

| countSelfReferences js1 == 0 = pure S.empty

| Function _ Nothing _ _ <- js1

, (argss, body, _) <- innerCollectAllFunctionArgs [] id js1

= S.insert (ident', length argss) <$> allInTailPosition body

| otherwise = empty

= guard (countSelfReferences js1 == 0) $> S.empty

= empty

toLoop :: S.Set Text -> Text -> Int -> [Text] -> [Text] -> AST -> State Int AST

toLoop trFns ident arity outerArgs innerArgs js = do

tcoDone <- tcoDoneM

modify (+ 1)

let

markDone :: Maybe SourceSpan -> AST

markDone ss = Assignment ss (Var ss tcoDone) (BooleanLiteral ss True)

loopify :: AST -> AST

loopify (Return ss ret)

| isSelfCall ident arity ret =

let

allArgumentValues = concat $ collectArgs [] ret

in

Block ss $

zipWith (\val arg ->

Assignment ss (Var ss (tcoVar arg)) val) allArgumentValues outerArgs

++ zipWith (\val arg ->

Assignment ss (Var ss (copyVar arg)) val) (drop (length outerArgs) allArgumentValues) innerArgs

++ [ ReturnNoResult ss ]

| isIndirectSelfCall ret = Return ss ret

| otherwise = Block ss [ markDone ss, Return ss ret ]

loopify (ReturnNoResult ss) = Block ss [ markDone ss, ReturnNoResult ss ]

loopify (While ss cond body) = While ss cond (loopify body)

loopify (For ss i js1 js2 body) = For ss i js1 js2 (loopify body)

loopify (ForIn ss i js1 body) = ForIn ss i js1 (loopify body)

loopify (IfElse ss cond body el) = IfElse ss cond (loopify body) (fmap loopify el)

loopify (Block ss body) = Block ss (map loopify body)

loopify (VariableIntroduction ss f (Just fn@(Function _ Nothing _ _)))

| (_, body, replace) <- innerCollectAllFunctionArgs [] id fn

, f `S.member` trFns = VariableIntroduction ss f (Just (replace (loopify body)))

loopify other = other

pure $ Block rootSS $

isIndirectSelfCall :: AST -> Bool

isIndirectSelfCall (App _ (Var _ ident') _) = ident' `S.member` trFns

isIndirectSelfCall (App _ fn _) = isIndirectSelfCall fn

isIndirectSelfCall _ = False

isSelfCall :: Text -> Int -> AST -> Bool

isSelfCall ident 1 (App _ (Var _ ident') _) = ident == ident'

isSelfCall ident arity (App _ fn _) = isSelfCall ident (arity - 1) fn

isSelfCall _ _ _ = False

foldMapA :: (Applicative f, Monoid w, Foldable t) => (a -> f w) -> t a -> f w

foldMapA f = foldr (liftA2 mappend . f) (pure mempty)

module Main where

import Prelude

import Effect.Console (log)

import Test.Assert (assertEqual)

data Maybe a = Nothing | Just a

f :: Int -> Int

f x = case x of

0 -> 0

n | _ <- n -> f (x - 1)

_ -> f (x - 2)

g :: Int -> Int

g x = case x of

0 -> 0

n | n == n, true -> g (x - 1)

_ -> g (x - 2)

weirdsum :: Int -> (Int -> Maybe Int) -> Int -> Int

weirdsum accum f n = case n of

0 -> accum

x | Just y <- f x -> weirdsum (accum + y) f (n - 1)

_ -> weirdsum accum f (n - 1)

tricksyinners :: Int -> Int -> Int

tricksyinners accum x = case x of

0 -> accum + f' x * f' x

n -> tricksyinners (accum + 2) (n - 1)

f' y = y + 3

main = do

assertEqual { expected: 0, actual: f 100000 }

assertEqual { expected: 0, actual: g 100000 }

assertEqual { expected: 20, actual: weirdsum 0 (\x -> if x < 5 then Just (2 * x) else Nothing) 100000 }

assertEqual { expected: 200009, actual: tricksyinners 0 100000 }

log "Done"