{-# LANGUAGE ViewPatterns #-}

module Refact.Fixity (applyFixities) where

import Control.Monad.Trans.State
import Data.Generics hiding (Fixity)
import Data.Maybe
import qualified GHC
import Language.Haskell.GHC.ExactPrint
import Refact.Compat (Fixity (..), SourceText (..), occNameString, rdrNameOcc)
import Refact.Utils

-- | Rearrange infix expressions to account for fixity.
-- The set of fixities is wired in and includes all fixities in base.
applyFixities :: Module -> IO Module
applyFixities :: Module -> IO Module
applyFixities Module
m = forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (forall (m :: * -> *). Monad m => GenericM m -> GenericM m
everywhereM (forall (m :: * -> *) a b.
(Monad m, Typeable a, Typeable b) =>
(b -> m b) -> a -> m a
mkM Expr -> StateT () IO Expr
expFix) Module
m) ()
  -- Note: everywhereM is a bottom-up transformation

expFix :: Expr -> StateT () IO Expr
expFix :: Expr -> StateT () IO Expr
expFix (GHC.L SrcSpanAnnA
loc (GHC.OpApp XOpApp GhcPs
an Expr
l Expr
op Expr
r)) =
  [(String, Fixity)]
-> SrcSpanAnnA
-> EpAnn [AddEpAnn]
-> Expr
-> Expr
-> Fixity
-> Expr
-> StateT () IO Expr
mkOpAppRn [(String, Fixity)]
baseFixities SrcSpanAnnA
loc XOpApp GhcPs
an Expr
l Expr
op ([(String, Fixity)] -> Expr -> Fixity
findFixity [(String, Fixity)]
baseFixities Expr
op) Expr
r
expFix Expr
e = forall (m :: * -> *) a. Monad m => a -> m a
return Expr
e

getIdent :: Expr -> String
getIdent :: Expr -> String
getIdent (forall l e. GenLocated l e -> e
GHC.unLoc -> GHC.HsVar XVar GhcPs
_ (GHC.L SrcSpanAnnN
_ RdrName
n)) = OccName -> String
occNameString forall b c a. (b -> c) -> (a -> b) -> a -> c
. RdrName -> OccName
rdrNameOcc forall a b. (a -> b) -> a -> b
$ RdrName
n
getIdent Expr
_ = forall a. HasCallStack => String -> a
error String
"Must be HsVar"

---------------------------
-- Modified from GHC Renamer
mkOpAppRn ::
  [(String, GHC.Fixity)] ->
  GHC.SrcSpanAnnA ->
  GHC.EpAnn [GHC.AddEpAnn] ->
  Expr -> -- Left operand; already rearranged
  Expr ->
  GHC.Fixity -> -- Operator and fixity
  Expr -> -- Right operand (not an OpApp, but might
  -- be a NegApp)
  StateT () IO Expr
-- (e11 `op1` e12) `op2` e2
mkOpAppRn :: [(String, Fixity)]
-> SrcSpanAnnA
-> EpAnn [AddEpAnn]
-> Expr
-> Expr
-> Fixity
-> Expr
-> StateT () IO Expr
mkOpAppRn [(String, Fixity)]
fs SrcSpanAnnA
loc EpAnn [AddEpAnn]
an e1 :: Expr
e1@(GHC.L SrcSpanAnnA
_ (GHC.OpApp XOpApp GhcPs
x1 Expr
e11 Expr
op1 Expr
e12)) Expr
op2 Fixity
fix2 Expr
e2
  | Bool
nofix_error =
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall l e. l -> e -> GenLocated l e
GHC.L SrcSpanAnnA
loc (forall p.
XOpApp p -> LHsExpr p -> LHsExpr p -> LHsExpr p -> HsExpr p
GHC.OpApp EpAnn [AddEpAnn]
an Expr
e1 Expr
op2 Expr
e2)
  | Bool
associate_right = do
    -- liftIO $ putStrLn $ "mkOpAppRn:1:e1" ++ showAst e1
    let e12' :: GenLocated SrcSpanAnnA (HsExpr GhcPs)
e12' = forall t a. Default t => LocatedAn t a -> DeltaPos -> LocatedAn t a
setEntryDP Expr
e12 (Int -> DeltaPos
GHC.SameLine Int
0)
    GenLocated SrcSpanAnnA (HsExpr GhcPs)
new_e <- [(String, Fixity)]
-> SrcSpanAnnA
-> EpAnn [AddEpAnn]
-> Expr
-> Expr
-> Fixity
-> Expr
-> StateT () IO Expr
mkOpAppRn [(String, Fixity)]
fs SrcSpanAnnA
loc' EpAnn [AddEpAnn]
an GenLocated SrcSpanAnnA (HsExpr GhcPs)
e12' Expr
op2 Fixity
fix2 Expr
e2
    let (GenLocated SrcSpanAnnA (HsExpr GhcPs)
new_e',Int
_,[String]
_) = forall a. Transform a -> (a, Int, [String])
runTransform forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) t2 t1 a b.
(Monad m, Monoid t2, Typeable t1, Typeable t2) =>
LocatedAn t1 a -> LocatedAn t2 b -> TransformT m (LocatedAn t2 b)
transferEntryDP Expr
e12 GenLocated SrcSpanAnnA (HsExpr GhcPs)
new_e
    let res :: GenLocated SrcSpanAnnA (HsExpr GhcPs)
res = forall l e. l -> e -> GenLocated l e
GHC.L SrcSpanAnnA
loc (forall p.
XOpApp p -> LHsExpr p -> LHsExpr p -> LHsExpr p -> HsExpr p
GHC.OpApp XOpApp GhcPs
x1 Expr
e11 Expr
op1 GenLocated SrcSpanAnnA (HsExpr GhcPs)
new_e')
    -- liftIO $ putStrLn $ "mkOpAppRn:1:res" ++ showAst res
    forall (m :: * -> *) a. Monad m => a -> m a
return GenLocated SrcSpanAnnA (HsExpr GhcPs)
res
  where
    loc' :: SrcSpanAnnA
loc' = forall a e1 e2.
Semigroup a =>
GenLocated (SrcAnn a) e1 -> GenLocated (SrcAnn a) e2 -> SrcAnn a
GHC.combineLocsA Expr
e12 Expr
e2
    fix1 :: Fixity
fix1 = [(String, Fixity)] -> Expr -> Fixity
findFixity [(String, Fixity)]
fs Expr
op1
    (Bool
nofix_error, Bool
associate_right) = Fixity -> Fixity -> (Bool, Bool)
GHC.compareFixity Fixity
fix1 Fixity
fix2

---------------------------
--      (- neg_arg) `op` e2
mkOpAppRn [(String, Fixity)]
fs SrcSpanAnnA
loc EpAnn [AddEpAnn]
an e1 :: Expr
e1@(GHC.L SrcSpanAnnA
_ (GHC.NegApp XNegApp GhcPs
an' Expr
neg_arg SyntaxExpr GhcPs
neg_name)) Expr
op2 Fixity
fix2 Expr
e2
  | Bool
nofix_error =
    forall (m :: * -> *) a. Monad m => a -> m a
return (forall l e. l -> e -> GenLocated l e
GHC.L SrcSpanAnnA
loc (forall p.
XOpApp p -> LHsExpr p -> LHsExpr p -> LHsExpr p -> HsExpr p
GHC.OpApp EpAnn [AddEpAnn]
an Expr
e1 Expr
op2 Expr
e2))
  | Bool
associate_right =
    do
      -- liftIO $ putStrLn $ "mkOpAppRn:2:e1" ++ showAst e1
      GenLocated SrcSpanAnnA (HsExpr GhcPs)
new_e <- [(String, Fixity)]
-> SrcSpanAnnA
-> EpAnn [AddEpAnn]
-> Expr
-> Expr
-> Fixity
-> Expr
-> StateT () IO Expr
mkOpAppRn [(String, Fixity)]
fs SrcSpanAnnA
loc' EpAnn [AddEpAnn]
an Expr
neg_arg Expr
op2 Fixity
fix2 Expr
e2
      let new_e' :: GenLocated SrcSpanAnnA (HsExpr GhcPs)
new_e' = forall t a. Default t => LocatedAn t a -> DeltaPos -> LocatedAn t a
setEntryDP GenLocated SrcSpanAnnA (HsExpr GhcPs)
new_e (Int -> DeltaPos
GHC.SameLine Int
0)
      let res :: GenLocated SrcSpanAnnA (HsExpr GhcPs)
res = forall t a. Default t => LocatedAn t a -> DeltaPos -> LocatedAn t a
setEntryDP (forall l e. l -> e -> GenLocated l e
GHC.L SrcSpanAnnA
loc (forall p. XNegApp p -> LHsExpr p -> SyntaxExpr p -> HsExpr p
GHC.NegApp XNegApp GhcPs
an' GenLocated SrcSpanAnnA (HsExpr GhcPs)
new_e' SyntaxExpr GhcPs
neg_name)) (Int -> DeltaPos
GHC.SameLine Int
0)
      -- liftIO $ putStrLn $ "mkOpAppRn:2:res" ++ showAst res
      forall (m :: * -> *) a. Monad m => a -> m a
return GenLocated SrcSpanAnnA (HsExpr GhcPs)
res
  where
    loc' :: SrcSpanAnnA
loc' = forall a e1 e2.
Semigroup a =>
GenLocated (SrcAnn a) e1 -> GenLocated (SrcAnn a) e2 -> SrcAnn a
GHC.combineLocsA Expr
neg_arg Expr
e2
    (Bool
nofix_error, Bool
associate_right) = Fixity -> Fixity -> (Bool, Bool)
GHC.compareFixity Fixity
GHC.negateFixity Fixity
fix2

---------------------------
--      e1 `op` - neg_arg
mkOpAppRn [(String, Fixity)]
_ SrcSpanAnnA
loc EpAnn [AddEpAnn]
an Expr
e1 Expr
op1 Fixity
fix1 e2 :: Expr
e2@(GHC.L SrcSpanAnnA
_ GHC.NegApp {}) -- NegApp can occur on the right
  | Bool -> Bool
not Bool
associate_right -- We *want* right association
    = do
    -- liftIO $ putStrLn $ "mkOpAppRn:3:e1" ++ showAst (GHC.L loc (GHC.OpApp an e1 op1 e2))
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall l e. l -> e -> GenLocated l e
GHC.L SrcSpanAnnA
loc (forall p.
XOpApp p -> LHsExpr p -> LHsExpr p -> LHsExpr p -> HsExpr p
GHC.OpApp EpAnn [AddEpAnn]
an Expr
e1 Expr
op1 Expr
e2)
  where
    (Bool
_, Bool
associate_right) = Fixity -> Fixity -> (Bool, Bool)
GHC.compareFixity Fixity
fix1 Fixity
GHC.negateFixity

---------------------------
--      Default case
mkOpAppRn [(String, Fixity)]
_ SrcSpanAnnA
loc EpAnn [AddEpAnn]
an Expr
e1 Expr
op Fixity
_fix Expr
e2 -- Default case, no rearrangment
  = do
  -- liftIO $ putStrLn $ "mkOpAppRn:4:e1" ++ showAst (GHC.L loc (GHC.OpApp an e1 op e2))
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall l e. l -> e -> GenLocated l e
GHC.L SrcSpanAnnA
loc (forall p.
XOpApp p -> LHsExpr p -> LHsExpr p -> LHsExpr p -> HsExpr p
GHC.OpApp EpAnn [AddEpAnn]
an Expr
e1 Expr
op Expr
e2)

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

findFixity :: [(String, GHC.Fixity)] -> Expr -> GHC.Fixity
findFixity :: [(String, Fixity)] -> Expr -> Fixity
findFixity [(String, Fixity)]
fs Expr
r = [(String, Fixity)] -> String -> Fixity
askFix [(String, Fixity)]
fs (Expr -> String
getIdent Expr
r)

askFix :: [(String, GHC.Fixity)] -> String -> GHC.Fixity
askFix :: [(String, Fixity)] -> String -> Fixity
askFix [(String, Fixity)]
xs = \String
k -> forall {a} {a}. Eq a => a -> a -> [(a, a)] -> a
lookupWithDefault Fixity
GHC.defaultFixity String
k [(String, Fixity)]
xs
  where
    lookupWithDefault :: a -> a -> [(a, a)] -> a
lookupWithDefault a
def_v a
k [(a, a)]
mp1 = forall a. a -> Maybe a -> a
fromMaybe a
def_v forall a b. (a -> b) -> a -> b
$ forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup a
k [(a, a)]
mp1

-- | All fixities defined in the Prelude.
preludeFixities :: [(String, GHC.Fixity)]
preludeFixities :: [(String, Fixity)]
preludeFixities =
  forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
    [ Int -> [String] -> [(String, Fixity)]
infixr_ Int
9 [String
"."],
      Int -> [String] -> [(String, Fixity)]
infixl_ Int
9 [String
"!!"],
      Int -> [String] -> [(String, Fixity)]
infixr_ Int
8 [String
"^", String
"^^", String
"**"],
      Int -> [String] -> [(String, Fixity)]
infixl_ Int
7 [String
"*", String
"/", String
"quot", String
"rem", String
"div", String
"mod", String
":%", String
"%"],
      Int -> [String] -> [(String, Fixity)]
infixl_ Int
6 [String
"+", String
"-"],
      Int -> [String] -> [(String, Fixity)]
infixr_ Int
5 [String
":", String
"++"],
      Int -> [String] -> [(String, Fixity)]
infix_ Int
4 [String
"==", String
"/=", String
"<", String
"<=", String
">=", String
">", String
"elem", String
"notElem"],
      Int -> [String] -> [(String, Fixity)]
infixr_ Int
3 [String
"&&"],
      Int -> [String] -> [(String, Fixity)]
infixr_ Int
2 [String
"||"],
      Int -> [String] -> [(String, Fixity)]
infixl_ Int
1 [String
">>", String
">>="],
      Int -> [String] -> [(String, Fixity)]
infixr_ Int
1 [String
"=<<"],
      Int -> [String] -> [(String, Fixity)]
infixr_ Int
0 [String
"$", String
"$!", String
"seq"]
    ]

-- | All fixities defined in the base package.
--
--   Note that the @+++@ operator appears in both Control.Arrows and
--   Text.ParserCombinators.ReadP. The listed precedence for @+++@ in
--   this list is that of Control.Arrows.
baseFixities :: [(String, GHC.Fixity)]
baseFixities :: [(String, Fixity)]
baseFixities =
  [(String, Fixity)]
preludeFixities
    forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
      [ Int -> [String] -> [(String, Fixity)]
infixl_ Int
9 [String
"!", String
"//", String
"!:"],
        Int -> [String] -> [(String, Fixity)]
infixl_ Int
8 [String
"shift", String
"rotate", String
"shiftL", String
"shiftR", String
"rotateL", String
"rotateR"],
        Int -> [String] -> [(String, Fixity)]
infixl_ Int
7 [String
".&."],
        Int -> [String] -> [(String, Fixity)]
infixl_ Int
6 [String
"xor"],
        Int -> [String] -> [(String, Fixity)]
infix_ Int
6 [String
":+"],
        Int -> [String] -> [(String, Fixity)]
infixl_ Int
5 [String
".|."],
        Int -> [String] -> [(String, Fixity)]
infixr_ Int
5 [String
"+:+", String
"<++", String
"<+>"], -- fixity conflict for +++ between ReadP and Arrow
        Int -> [String] -> [(String, Fixity)]
infix_ Int
5 [String
"\\\\"],
        Int -> [String] -> [(String, Fixity)]
infixl_ Int
4 [String
"<$>", String
"<$", String
"<*>", String
"<*", String
"*>", String
"<**>"],
        Int -> [String] -> [(String, Fixity)]
infix_ Int
4 [String
"elemP", String
"notElemP"],
        Int -> [String] -> [(String, Fixity)]
infixl_ Int
3 [String
"<|>"],
        Int -> [String] -> [(String, Fixity)]
infixr_ Int
3 [String
"&&&", String
"***"],
        Int -> [String] -> [(String, Fixity)]
infixr_ Int
2 [String
"+++", String
"|||"],
        Int -> [String] -> [(String, Fixity)]
infixr_ Int
1 [String
"<=<", String
">=>", String
">>>", String
"<<<", String
"^<<", String
"<<^", String
"^>>", String
">>^"],
        Int -> [String] -> [(String, Fixity)]
infixl_ Int
0 [String
"on"],
        Int -> [String] -> [(String, Fixity)]
infixr_ Int
0 [String
"par", String
"pseq"]
      ]

infixr_, infixl_, infix_ :: Int -> [String] -> [(String, GHC.Fixity)]
infixr_ :: Int -> [String] -> [(String, Fixity)]
infixr_ = FixityDirection -> Int -> [String] -> [(String, Fixity)]
fixity FixityDirection
GHC.InfixR
infixl_ :: Int -> [String] -> [(String, Fixity)]
infixl_ = FixityDirection -> Int -> [String] -> [(String, Fixity)]
fixity FixityDirection
GHC.InfixL
infix_ :: Int -> [String] -> [(String, Fixity)]
infix_ = FixityDirection -> Int -> [String] -> [(String, Fixity)]
fixity FixityDirection
GHC.InfixN

-- Internal: help function for the above definitions.
fixity :: GHC.FixityDirection -> Int -> [String] -> [(String, GHC.Fixity)]
fixity :: FixityDirection -> Int -> [String] -> [(String, Fixity)]
fixity FixityDirection
a Int
p = forall a b. (a -> b) -> [a] -> [b]
map (,SourceText -> Int -> FixityDirection -> Fixity
Fixity (String -> SourceText
SourceText String
"") Int
p FixityDirection
a)