@@ -15,6 +15,8 @@
1515
1616module Language.PureScript.CodeGen.JS.Optimizer.TCO (tco) where
1717
18+import Data.Monoid
19+
1820import Language.PureScript.Options
1921import Language.PureScript.CodeGen.JS.AST
2022
@@ -30,10 +32,13 @@ tco' = everywhereOnJS convert
3032 where
3133 tcoLabel :: String
3234 tcoLabel = "tco"
35+
3336 tcoVar :: String -> String
3437 tcoVar arg = "__tco_" ++ arg
38+
3539 copyVar :: String -> String
3640 copyVar arg = "__copy_" ++ arg
41+
3742 convert :: JS -> JS
3843 convert js@(JSVariableIntroduction name (Just fn@JSFunction {})) =
3944 let
@@ -46,6 +51,7 @@ tco' = everywhereOnJS convert
4651 JSVariableIntroduction name (Just (replace (toLoop name allArgs body')))
4752 | otherwise -> js
4853 convert js = js
54+
4955 collectAllFunctionArgs :: [[String]] -> (JS -> JS) -> JS -> ([[String]], JS, JS -> JS)
5056 collectAllFunctionArgs allArgs f (JSFunction ident args (JSBlock (body@(JSReturn _):_))) =
5157 collectAllFunctionArgs (args : allArgs) (\b -> f (JSFunction ident (map copyVar args) (JSBlock [b]))) body
@@ -56,25 +62,35 @@ tco' = everywhereOnJS convert
5662 collectAllFunctionArgs allArgs f (JSReturn (JSFunction ident args body@(JSBlock _))) =
5763 (args : allArgs, body, f . JSReturn . JSFunction ident (map copyVar args))
5864 collectAllFunctionArgs allArgs f body = (allArgs, body, f)
65+
5966 isTailCall :: String -> JS -> Bool
6067 isTailCall ident js =
6168 let
6269 numSelfCalls = everythingOnJS (+) countSelfCalls js
6370 numSelfCallsInTailPosition = everythingOnJS (+) countSelfCallsInTailPosition js
6471 numSelfCallsUnderFunctions = everythingOnJS (+) countSelfCallsUnderFunctions js
72+ numSelfCallWithFnArgs = everythingOnJS (+) countSelfCallsWithFnArgs js
6573 in
6674 numSelfCalls > 0
6775 && numSelfCalls == numSelfCallsInTailPosition
6876 && numSelfCallsUnderFunctions == 0
77+ && numSelfCallWithFnArgs == 0
6978 where
7079 countSelfCalls :: JS -> Int
7180 countSelfCalls (JSApp (JSVar ident') _) | ident == ident' = 1
7281 countSelfCalls _ = 0
82+
7383 countSelfCallsInTailPosition :: JS -> Int
7484 countSelfCallsInTailPosition (JSReturn ret) | isSelfCall ident ret = 1
7585 countSelfCallsInTailPosition _ = 0
86+
87+ countSelfCallsUnderFunctions :: JS -> Int
7688 countSelfCallsUnderFunctions (JSFunction _ _ js') = everythingOnJS (+) countSelfCalls js'
7789 countSelfCallsUnderFunctions _ = 0
90+
91+ countSelfCallsWithFnArgs :: JS -> Int
92+ countSelfCallsWithFnArgs ret = if isSelfCallWithFnArgs ident ret [] then 1 else 0
93+
7894 toLoop :: String -> [String] -> JS -> JS
7995 toLoop ident allArgs js = JSBlock $
8096 map (\arg -> JSVariableIntroduction arg (Just (JSVar (copyVar arg)))) allArgs ++
@@ -94,10 +110,19 @@ tco' = everywhereOnJS convert
94110 collectSelfCallArgs :: [[JS]] -> JS -> [[JS]]
95111 collectSelfCallArgs allArgumentValues (JSApp fn args') = collectSelfCallArgs (args' : allArgumentValues) fn
96112 collectSelfCallArgs allArgumentValues _ = allArgumentValues
113+
97114 isSelfCall :: String -> JS -> Bool
98− isSelfCall ident (JSApp (JSVar ident') args) | ident == ident' && not (any isFunction args) = True
99− isSelfCall ident (JSApp fn args) | not (any isFunction args) = isSelfCall ident fn
115+ isSelfCall ident (JSApp (JSVar ident') _) = ident == ident'
116+ isSelfCall ident (JSApp fn _) = isSelfCall ident fn
100117 isSelfCall _ _ = False
101− isFunction :: JS -> Bool
102− isFunction (JSFunction _ _ _) = True
103− isFunction _ = False
118+
119+ isSelfCallWithFnArgs :: String -> JS -> [JS] -> Bool
120+ isSelfCallWithFnArgs ident (JSVar ident') args | ident == ident' && any hasFunction args = True
121+ isSelfCallWithFnArgs ident (JSApp fn args) acc = isSelfCallWithFnArgs ident fn (args ++ acc)
122+ isSelfCallWithFnArgs _ _ _ = False
123+
124+ hasFunction :: JS -> Bool
125+ hasFunction = getAny . everythingOnJS mappend (Any . isFunction)
126+ where
127+ isFunction (JSFunction _ _ _) = True
128+ isFunction _ = False