OptimizeTailRecursion.hs

{-# language OverloadedStrings #-}
{-# language DataKinds #-}
{-# language BangPatterns #-}
module OptimizeTailRecursion where

import Control.Applicative ((<|>))
import Control.Lens.Cons (_last, _init)
import Control.Lens.Fold ((^..), (^?), (^?!), allOf, anyOf, folded, foldrOf)
import Control.Lens.Getter ((^.), to)
import Control.Lens.Plated (cosmos, transform, transformOn)
import Control.Lens.Prism (_Just)
import Control.Lens.Review ((#))
import Control.Lens.Setter ((%~), (.~))
import Control.Lens.Tuple (_2, _3)
import Data.Foldable (toList)
import Data.Function ((&))
import Data.Semigroup ((<>))

import Language.Python.Optics
import Language.Python.DSL
import Language.Python.Syntax.Expr (Expr (..), _Exprs, argExpr, paramName)
import Language.Python.Syntax.Statement (CompoundStatement (..), Statement (..), SmallStatement (..), SimpleStatement (..), _Statements)

optimizeTailRecursion :: Raw Statement -> Maybe (Raw Statement)
optimizeTailRecursion st = do
  function <- st ^? _Fundef
  let functionBody = function ^. body_
  bodyLast <- lastStatement functionBody

  let
    functionName = function ^. fdName.identValue
    bodyInit = functionBody ^?! _init
    paramNames = function ^.. fdParameters.folded.paramName.identValue

  if not $ hasTC functionName bodyLast
    then Nothing
    else
      Just $
      _Fundef #
        (function &
         body_ .~
           (zipWith
              (\a b -> line_ (var_ (a <> "__tr") .= var_ b))
              paramNames
              paramNames <>

            [ line_ ("__res__tr" .= none_)
            , line_ . while_ true_ .
              transformOn (traverse._Exprs) (renameIn paramNames "__tr") $
                bodyInit <>
                looped functionName paramNames bodyLast
            , line_ $ return_ "__res__tr"
            ]))

  where
    lastStatement :: [Raw Line] -> Maybe (Raw Statement)
    lastStatement = go Nothing
      where
        go !res [] = res
        go !res (a:as) = go (a ^? _Statements <|> res) as

    isTailCall :: String -> Raw Expr -> Bool
    isTailCall name e
      | anyOf (cosmos._Call.callFunction._Ident.identValue) (== name) e
      = (e ^? _Call.callFunction._Ident.identValue) == Just name
      | otherwise = False

    hasTC :: String -> Raw Statement -> Bool
    hasTC name st =
      case st of
        CompoundStatement (If _ _ _ _ sts [] sts') ->
          allOf _last (hasTC name) (sts ^.. _Statements) ||
          allOf _last (hasTC name) (sts' ^.. _Just._3._Statements)
        SmallStatement _ (MkSmallStatement s ss _ _ _) ->
          case last (s : fmap (^. _2) ss) of
            Return _ _ (Just e) -> isTailCall name e
            -- Return _ _ Nothing -> True
            Expr _ e -> isTailCall name e
            _ -> False
        _ -> False

    renameIn :: [String] -> String -> Raw Expr -> Raw Expr
    renameIn params suffix =
      transform
        (_Ident.identValue %~ (\a -> if a `elem` params then a <> suffix else a))

    looped :: String -> [String] -> Raw Statement -> [Raw Line]
    looped name params st
      | Just ifSt <- st ^? _If
      , hasTC name st =
          let
            ifBodyLines = toList $ ifSt ^. body_
          in
            case ifSt ^? to getElse._Just.body_ of
              Nothing ->
                [ line_ $
                  if_ (ifSt ^. ifCond)
                    ((ifBodyLines ^?! _init) <>
                     looped name params (ifBodyLines ^?! _last._Statements))
                ]
              Just sts'' ->
                [ line_ $
                  if_ (ifSt ^. ifCond)
                    ((ifSt ^?! body_.to toList._init) <>
                     looped name params (ifBodyLines ^?! _last._Statements)) &
                  else_
                    ((toList sts'' ^?! _init) <>
                     looped name params (toList sts'' ^?! _last._Statements))
                ]
      | otherwise =
          case st of
            CompoundStatement{} -> [line_ st]
            SmallStatement idnts (MkSmallStatement s ss sc cmt nl) ->
              let
                initExps = foldr (\_ _ -> init ss) [] ss
                lastExp = foldrOf (folded._2) (\_ _ -> last ss ^. _2) s ss
                newSts =
                  case initExps of
                    [] -> []
                    first : rest ->
                      [ line_ $
                        SmallStatement idnts
                        (MkSmallStatement (first ^. _2) rest sc cmt nl)
                      ]
              in
                case lastExp of
                  Return _ _ e ->
                    case e ^? _Just._Call of
                      Just call
                        | Just name' <- call ^? callFunction._Ident.identValue
                        , name' == name ->
                            newSts <>
                            fmap
                              (\a -> line_ (var_ (a <> "__tr__old") .= var_ (a <> "__tr")))
                              params <>
                            zipWith
                              (\a b -> line_ (var_ (a <> "__tr") .= b))
                              params
                              (transformOn
                                traverse
                                (renameIn params "__tr__old")
                                (call ^.. callArguments.folded.folded.argExpr))
                      _ ->
                        newSts <>
                        maybe [] (\e' -> [ line_ ("__res__tr" .= e') ]) e <>
                        [ line_ break_ ]
                  Expr _ e
                    | isTailCall name e -> newSts <> [line_ pass_]
                  _ -> [line_ st]