Function.hs

-- Copyright (c) 2025 Andrew Farmer
-- Copyright (c) 2020-2024 Facebook, Inc. and its affiliates.
--
-- This source code is licensed under the MIT license found in the
-- LICENSE file in the root directory of this source tree.
--
{-# LANGUAGE CPP #-}
{-# LANGUAGE TupleSections #-}
module Retrie.Rewrites.Function
  ( dfnsToRewrites
  , getImports
  , matchToRewrites
  ) where

import Control.Monad
import Control.Monad.State.Lazy
import Data.List
import Data.Maybe
import Data.Traversable

import Retrie.ExactPrint
import Retrie.Expr
import Retrie.GHC
import Retrie.Quantifiers
import Retrie.Types

dfnsToRewrites
  :: LibDir
  -> [(FastString, Direction)]
  -> AnnotatedModule
  -> IO (UniqFM FastString [Rewrite (LHsExpr GhcPs)])
dfnsToRewrites libdir specs am = fmap astA $ transformA am $ \ (L _ m) -> do
  let
    fsMap = uniqBag specs

  rrs <- sequence
    [ do
        fe <- mkLocatedHsVar fRdrName
        -- lift $ debugPrint Loud "dfnsToRewrites:ef="  [showAst fe]
        imps <- getImports libdir dir (hsmodName m)
        (fName,) . concat <$>
          forM (unLoc $ mg_alts $ fun_matches f) (matchToRewrites fe imps dir)
    | L _ (ValD _ f@FunBind{}) <- hsmodDecls m
    , let fRdrName = fun_id f
    , let fName = occNameFS (occName (unLoc fRdrName))
    , dir <- fromMaybe [] (lookupUFM fsMap fName)
    ]

  return $ listToUFM_C (++) rrs

------------------------------------------------------------------------

getImports
  :: LibDir -> Direction -> Maybe (LocatedA ModuleName) -> TransformT IO AnnotatedImports
getImports libdir RightToLeft (Just (L _ mn)) = -- See Note [fold only]
  TransformT $ lift $ liftIO $ parseImports libdir ["import " ++ moduleNameString mn]
getImports _ _ _ = return mempty

matchToRewrites
  :: LHsExpr GhcPs
  -> AnnotatedImports
  -> Direction
  -> LMatch GhcPs (LHsExpr GhcPs)
  -> TransformT IO [Rewrite (LHsExpr GhcPs)]
matchToRewrites e imps dir (L _ alt) = do
  let
    pats = getMatchPats alt
    grhss = m_grhss alt
  qss <- for (zip (inits pats) (tails pats)) $
    makeFunctionQuery e imps dir grhss mkApps
  qs <- backtickRules e imps dir grhss pats
  return $ qs ++ concat qss

type AppBuilder =
  LHsExpr GhcPs -> [LHsExpr GhcPs] -> TransformT IO (LHsExpr GhcPs)

irrefutablePat :: LPat GhcPs -> Bool
irrefutablePat = go . unLoc
  where
    go WildPat{} = True
    go VarPat{} = True
    go (LazyPat _ p) = irrefutablePat p
#if __GLASGOW_HASKELL__ < 912
    go (AsPat _ _ _ p) = irrefutablePat p
    go (ParPat _ _ p _) = irrefutablePat p
#else
    go (AsPat _ _ p) = irrefutablePat p
    go (ParPat _ p) = irrefutablePat p
#endif
    go (BangPat _ p) = irrefutablePat p
    go _ = False

makeFunctionQuery
  :: LHsExpr GhcPs
  -> AnnotatedImports
  -> Direction
  -> GRHSs GhcPs (LHsExpr GhcPs)
  -> AppBuilder
  -> ([LPat GhcPs], [LPat GhcPs])
  -> TransformT IO [Rewrite (LHsExpr GhcPs)]
makeFunctionQuery e imps dir grhss mkAppFn (argpats, bndpats)
  | any (not . irrefutablePat) bndpats = return []
  | otherwise = do
    let
      GRHSs _ _ lbs = grhss
      rhssList = grhssList grhss
      bs = collectPatsBinders CollNoDictBinders argpats
    -- See Note [Wildcards]
    (es,(_,bs')) <- runStateT (mapM patToExpr argpats) (wildSupply bs, bs)
    -- lift $ debugPrint Loud "makeFunctionQuery:e="  [showAst e]
    lhs <- mkAppFn e es
    for rhssList $ \ grhs -> do
      le <- mkLet lbs (grhsToExpr grhs)
      rhs <- mkLams bndpats le
      let
        (pat, temp) =
          case dir of
            LeftToRight -> (lhs,rhs)
            RightToLeft -> (rhs,lhs)
      p <- pruneA pat
      t <- pruneA temp
      return $ addRewriteImports imps $ mkRewrite (mkQs bs') p t

backtickRules
  :: LHsExpr GhcPs
  -> AnnotatedImports
  -> Direction
  -> GRHSs GhcPs (LHsExpr GhcPs)
  -> [LPat GhcPs]
  -> TransformT IO [Rewrite (LHsExpr GhcPs)]
backtickRules e imps dir@LeftToRight grhss (p1:p2:rest) = do
  let
#if __GLASGOW_HASKELL__ < 912
    na = noAnn
#else
    na = noExtField
#endif

    wrapOpApps op [] = pure op
    wrapOpApps op extra = do
      o <- mkParen op
      mkApps o extra

    both, left, right :: AppBuilder
    -- A function of arity greater than two used infix supplies its
    -- first two arguments via the operator and the rest by application.
    both op (l:r:extra) = do
      opApp <- mkLocA (SameLine 1) (OpApp na l op r)
      wrapOpApps opApp extra
    both _ _ = fail "backtickRules - both: impossible!"

    left op (l:extra) = do
      opSecL <- mkLocA (SameLine 1) (SectionL na l op)
      wrapOpApps opSecL extra
    left _ _ = fail "backtickRules - left: impossible!"

    right op (r:extra) = do
      opSecR <- mkLocA (SameLine 1) (SectionR na op r)
      wrapOpApps opSecR extra
    right _ _ = fail "backtickRules - right: impossible!"

    splits r = for (zip (inits r) (tails r))

  -- (p1 `op` p2) rest
  qss <- splits rest $ \(ri, rt) ->
    makeFunctionQuery e imps dir grhss both (p1 : p2 : ri, rt)
  -- (p1 `op`) p2 rest
  qsl <- splits (p2:rest) $ \(ri, rt) ->
    makeFunctionQuery e imps dir grhss left (p1 : ri, rt)
  -- (`op` p2) p1 rest
  qsr <- splits (p1:rest) $ \(ri, rt) ->
    makeFunctionQuery e imps dir grhss right (p2 : ri, rt)
  return $ concat qss ++ concat qsl ++ concat qsr
backtickRules _ _ _ _ _ = return []

-- Note [fold only]
-- Currently we only generate imports for folds, because it is easy.
-- (We only need to add an import for the module defining the folded
-- function.) Generating the imports for unfolds will require some
-- sort of analysis with haskell-names and is a TODO.