{-# LANGUAGE RecordWildCards #-}
module Overloaded.Plugin.IdiomBrackets where

import Data.List          (foldl')
import Data.List.NonEmpty (NonEmpty (..))

import GHC.Compat.Expr

import Overloaded.Plugin.Rewrite
import Overloaded.Plugin.Names

transformIdiomBrackets
    :: Names
    -> LHsExpr GhcRn
    -> Rewrite (LHsExpr GhcRn)
transformIdiomBrackets :: Names -> LHsExpr GhcRn -> Rewrite (LHsExpr GhcRn)
transformIdiomBrackets Names
names (L SrcSpan
_l (HsRnBracketOut XRnBracketOut GhcRn
_ (ExpBr XExpBr GhcRn
_ LHsExpr GhcRn
e) [PendingRnSplice]
_))
    = LHsExpr GhcRn -> Rewrite (LHsExpr GhcRn)
forall a. a -> Rewrite a
Rewrite (Names -> LHsExpr GhcRn -> LHsExpr GhcRn
transformIdiomBrackets' Names
names LHsExpr GhcRn
e)
transformIdiomBrackets Names
_ LHsExpr GhcRn
_ = Rewrite (LHsExpr GhcRn)
forall a. Rewrite a
NoRewrite

transformIdiomBrackets'
    :: Names
    -> LHsExpr GhcRn
    -> LHsExpr GhcRn
transformIdiomBrackets' :: Names -> LHsExpr GhcRn -> LHsExpr GhcRn
transformIdiomBrackets' Names
names expr :: LHsExpr GhcRn
expr@(L SrcSpan
_e OpApp {}) = do
    let bt :: BT (LHsExpr GhcRn)
bt = LHsExpr GhcRn -> BT (LHsExpr GhcRn)
forall p. LHsExpr p -> BT (LHsExpr p)
matchOp LHsExpr GhcRn
expr
    let result :: LHsExpr GhcRn
result = Names -> BT (LHsExpr GhcRn) -> LHsExpr GhcRn
idiomBT Names
names BT (LHsExpr GhcRn)
bt
    LHsExpr GhcRn
result
transformIdiomBrackets' Names
names LHsExpr GhcRn
expr = do
    let (LHsExpr GhcRn
f :| [LHsExpr GhcRn]
args) = LHsExpr GhcRn -> NonEmpty (LHsExpr GhcRn)
forall p. LHsExpr p -> NonEmpty (LHsExpr p)
matchApp LHsExpr GhcRn
expr
    let f' :: LHsExpr GhcRn
f' = Names -> LHsExpr GhcRn -> LHsExpr GhcRn
pureExpr Names
names LHsExpr GhcRn
f
    let result :: LHsExpr GhcRn
result = (LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn)
-> LHsExpr GhcRn -> [LHsExpr GhcRn] -> LHsExpr GhcRn
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
applyExpr Names
names) LHsExpr GhcRn
f' [LHsExpr GhcRn]
args
    LHsExpr GhcRn
result

-------------------------------------------------------------------------------
-- Function application maching
-------------------------------------------------------------------------------

-- | Match nested function applications, 'HsApp':
-- f x y z ~> f :| [x,y,z]
--
matchApp :: LHsExpr p -> NonEmpty (LHsExpr p)
matchApp :: LHsExpr p -> NonEmpty (LHsExpr p)
matchApp (L SrcSpan
_ (HsApp XApp p
_ LHsExpr p
f LHsExpr p
x)) = NonEmpty (LHsExpr p) -> LHsExpr p -> NonEmpty (LHsExpr p)
forall a. NonEmpty a -> a -> NonEmpty a
neSnoc (LHsExpr p -> NonEmpty (LHsExpr p)
forall p. LHsExpr p -> NonEmpty (LHsExpr p)
matchApp LHsExpr p
f) LHsExpr p
x
matchApp LHsExpr p
e = LHsExpr p -> NonEmpty (LHsExpr p)
forall (f :: * -> *) a. Applicative f => a -> f a
pure LHsExpr p
e

neSnoc :: NonEmpty a -> a -> NonEmpty a
neSnoc :: NonEmpty a -> a -> NonEmpty a
neSnoc (a
x :| [a]
xs) a
y = a
x a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
y]

-------------------------------------------------------------------------------
-- Operator application matching
-------------------------------------------------------------------------------

-- | Match nested operator applications, 'OpApp'.
-- x + y * z ~>  Branch (+) (Leaf x) (Branch (*) (Leaf y) (Leaf z))
matchOp :: LHsExpr p -> BT (LHsExpr p)
matchOp :: LHsExpr p -> BT (LHsExpr p)
matchOp (L SrcSpan
_ (OpApp XOpApp p
_  LHsExpr p
lhs LHsExpr p
op LHsExpr p
rhs)) = BT (LHsExpr p) -> LHsExpr p -> BT (LHsExpr p) -> BT (LHsExpr p)
forall a. BT a -> a -> BT a -> BT a
Branch (LHsExpr p -> BT (LHsExpr p)
forall p. LHsExpr p -> BT (LHsExpr p)
matchOp LHsExpr p
lhs) LHsExpr p
op (LHsExpr p -> BT (LHsExpr p)
forall p. LHsExpr p -> BT (LHsExpr p)
matchOp LHsExpr p
rhs)
matchOp LHsExpr p
x = LHsExpr p -> BT (LHsExpr p)
forall a. a -> BT a
Leaf LHsExpr p
x

-- | Non-empty binary tree, with elements at branches too.
data BT a = Leaf a | Branch (BT a) a (BT a)

-- flatten: note that leaf is returned as is.
idiomBT :: Names -> BT (LHsExpr GhcRn) -> LHsExpr GhcRn
idiomBT :: Names -> BT (LHsExpr GhcRn) -> LHsExpr GhcRn
idiomBT Names
_     (Leaf LHsExpr GhcRn
x)            = LHsExpr GhcRn
x
idiomBT Names
names (Branch BT (LHsExpr GhcRn)
lhs LHsExpr GhcRn
op BT (LHsExpr GhcRn)
rhs) = Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
fmapExpr Names
names LHsExpr GhcRn
op (Names -> BT (LHsExpr GhcRn) -> LHsExpr GhcRn
idiomBT Names
names BT (LHsExpr GhcRn)
lhs) LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
`ap` Names -> BT (LHsExpr GhcRn) -> LHsExpr GhcRn
idiomBT Names
names BT (LHsExpr GhcRn)
rhs
  where
    ap :: LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
ap = Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
apExpr Names
names

-------------------------------------------------------------------------------
-- Idioms related constructors
-------------------------------------------------------------------------------

applyExpr :: Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
applyExpr :: Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
applyExpr Names
names LHsExpr GhcRn
f (L SrcSpan
_ (HsPar XPar GhcRn
_ (L SrcSpan
_ (HsApp XApp GhcRn
_ (L SrcSpan
_ (HsVar XVar GhcRn
_ (L SrcSpan
_ IdP GhcRn
voidName'))) LHsExpr GhcRn
x))))
    | IdP GhcRn
Name
voidName' Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Names -> Name
voidName Names
names = Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
birdExpr Names
names LHsExpr GhcRn
f LHsExpr GhcRn
x
applyExpr Names
names LHsExpr GhcRn
f LHsExpr GhcRn
x               = Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
apExpr Names
names LHsExpr GhcRn
f LHsExpr GhcRn
x

apExpr :: Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
apExpr :: Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
apExpr Names {Name
CatNames
catNames :: Names -> CatNames
codeFromStringName :: Names -> Name
codeFromLabelName :: Names -> Name
conRightName :: Names -> Name
conLeftName :: Names -> Name
doBindName :: Names -> Name
doThenName :: Names -> Name
doPureName :: Names -> Name
composeName :: Names -> Name
birdName :: Names -> Name
apName :: Names -> Name
pureName :: Names -> Name
fmapName :: Names -> Name
fromTypeSymbolName :: Names -> Name
fromTypeNatName :: Names -> Name
fromLabelName :: Names -> Name
unitName :: Names -> Name
ifteName :: Names -> Name
consName :: Names -> Name
nilName :: Names -> Name
fromCharName :: Names -> Name
fromNaturalName :: Names -> Name
fromNumeralName :: Names -> Name
fromSymbolName :: Names -> Name
fromStringName :: Names -> Name
catNames :: CatNames
codeFromStringName :: Name
codeFromLabelName :: Name
conRightName :: Name
conLeftName :: Name
doBindName :: Name
doThenName :: Name
doPureName :: Name
composeName :: Name
voidName :: Name
birdName :: Name
apName :: Name
pureName :: Name
fmapName :: Name
fromTypeSymbolName :: Name
fromTypeNatName :: Name
fromLabelName :: Name
unitName :: Name
ifteName :: Name
consName :: Name
nilName :: Name
fromCharName :: Name
fromNaturalName :: Name
fromNumeralName :: Name
fromSymbolName :: Name
fromStringName :: Name
voidName :: Names -> Name
..} LHsExpr GhcRn
f LHsExpr GhcRn
x = SrcSpan -> LHsExpr GhcRn -> [LHsExpr GhcRn] -> LHsExpr GhcRn
hsApps SrcSpan
l' (SrcSpan -> Name -> LHsExpr GhcRn
hsVar SrcSpan
l' Name
apName) [LHsExpr GhcRn
f, LHsExpr GhcRn
x] where
    l' :: SrcSpan
l' = SrcSpan
noSrcSpan

birdExpr :: Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
birdExpr :: Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
birdExpr Names {Name
CatNames
catNames :: CatNames
codeFromStringName :: Name
codeFromLabelName :: Name
conRightName :: Name
conLeftName :: Name
doBindName :: Name
doThenName :: Name
doPureName :: Name
composeName :: Name
voidName :: Name
birdName :: Name
apName :: Name
pureName :: Name
fmapName :: Name
fromTypeSymbolName :: Name
fromTypeNatName :: Name
fromLabelName :: Name
unitName :: Name
ifteName :: Name
consName :: Name
nilName :: Name
fromCharName :: Name
fromNaturalName :: Name
fromNumeralName :: Name
fromSymbolName :: Name
fromStringName :: Name
catNames :: Names -> CatNames
codeFromStringName :: Names -> Name
codeFromLabelName :: Names -> Name
conRightName :: Names -> Name
conLeftName :: Names -> Name
doBindName :: Names -> Name
doThenName :: Names -> Name
doPureName :: Names -> Name
composeName :: Names -> Name
birdName :: Names -> Name
apName :: Names -> Name
pureName :: Names -> Name
fmapName :: Names -> Name
fromTypeSymbolName :: Names -> Name
fromTypeNatName :: Names -> Name
fromLabelName :: Names -> Name
unitName :: Names -> Name
ifteName :: Names -> Name
consName :: Names -> Name
nilName :: Names -> Name
fromCharName :: Names -> Name
fromNaturalName :: Names -> Name
fromNumeralName :: Names -> Name
fromSymbolName :: Names -> Name
fromStringName :: Names -> Name
voidName :: Names -> Name
..} LHsExpr GhcRn
f LHsExpr GhcRn
x = SrcSpan -> LHsExpr GhcRn -> [LHsExpr GhcRn] -> LHsExpr GhcRn
hsApps SrcSpan
l' (SrcSpan -> Name -> LHsExpr GhcRn
hsVar SrcSpan
l' Name
birdName) [LHsExpr GhcRn
f, LHsExpr GhcRn
x] where
    l' :: SrcSpan
l' = SrcSpan
noSrcSpan

fmapExpr :: Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
fmapExpr :: Names -> LHsExpr GhcRn -> LHsExpr GhcRn -> LHsExpr GhcRn
fmapExpr Names {Name
CatNames
catNames :: CatNames
codeFromStringName :: Name
codeFromLabelName :: Name
conRightName :: Name
conLeftName :: Name
doBindName :: Name
doThenName :: Name
doPureName :: Name
composeName :: Name
voidName :: Name
birdName :: Name
apName :: Name
pureName :: Name
fmapName :: Name
fromTypeSymbolName :: Name
fromTypeNatName :: Name
fromLabelName :: Name
unitName :: Name
ifteName :: Name
consName :: Name
nilName :: Name
fromCharName :: Name
fromNaturalName :: Name
fromNumeralName :: Name
fromSymbolName :: Name
fromStringName :: Name
catNames :: Names -> CatNames
codeFromStringName :: Names -> Name
codeFromLabelName :: Names -> Name
conRightName :: Names -> Name
conLeftName :: Names -> Name
doBindName :: Names -> Name
doThenName :: Names -> Name
doPureName :: Names -> Name
composeName :: Names -> Name
birdName :: Names -> Name
apName :: Names -> Name
pureName :: Names -> Name
fmapName :: Names -> Name
fromTypeSymbolName :: Names -> Name
fromTypeNatName :: Names -> Name
fromLabelName :: Names -> Name
unitName :: Names -> Name
ifteName :: Names -> Name
consName :: Names -> Name
nilName :: Names -> Name
fromCharName :: Names -> Name
fromNaturalName :: Names -> Name
fromNumeralName :: Names -> Name
fromSymbolName :: Names -> Name
fromStringName :: Names -> Name
voidName :: Names -> Name
..} LHsExpr GhcRn
f LHsExpr GhcRn
x = SrcSpan -> LHsExpr GhcRn -> [LHsExpr GhcRn] -> LHsExpr GhcRn
hsApps SrcSpan
l' (SrcSpan -> Name -> LHsExpr GhcRn
hsVar SrcSpan
l' Name
fmapName) [LHsExpr GhcRn
f, LHsExpr GhcRn
x] where
    l' :: SrcSpan
l' = SrcSpan
noSrcSpan

pureExpr :: Names -> LHsExpr GhcRn -> LHsExpr GhcRn
pureExpr :: Names -> LHsExpr GhcRn -> LHsExpr GhcRn
pureExpr Names {Name
CatNames
catNames :: CatNames
codeFromStringName :: Name
codeFromLabelName :: Name
conRightName :: Name
conLeftName :: Name
doBindName :: Name
doThenName :: Name
doPureName :: Name
composeName :: Name
voidName :: Name
birdName :: Name
apName :: Name
pureName :: Name
fmapName :: Name
fromTypeSymbolName :: Name
fromTypeNatName :: Name
fromLabelName :: Name
unitName :: Name
ifteName :: Name
consName :: Name
nilName :: Name
fromCharName :: Name
fromNaturalName :: Name
fromNumeralName :: Name
fromSymbolName :: Name
fromStringName :: Name
catNames :: Names -> CatNames
codeFromStringName :: Names -> Name
codeFromLabelName :: Names -> Name
conRightName :: Names -> Name
conLeftName :: Names -> Name
doBindName :: Names -> Name
doThenName :: Names -> Name
doPureName :: Names -> Name
composeName :: Names -> Name
birdName :: Names -> Name
apName :: Names -> Name
pureName :: Names -> Name
fmapName :: Names -> Name
fromTypeSymbolName :: Names -> Name
fromTypeNatName :: Names -> Name
fromLabelName :: Names -> Name
unitName :: Names -> Name
ifteName :: Names -> Name
consName :: Names -> Name
nilName :: Names -> Name
fromCharName :: Names -> Name
fromNaturalName :: Names -> Name
fromNumeralName :: Names -> Name
fromSymbolName :: Names -> Name
fromStringName :: Names -> Name
voidName :: Names -> Name
..} LHsExpr GhcRn
x = SrcSpan -> LHsExpr GhcRn -> [LHsExpr GhcRn] -> LHsExpr GhcRn
hsApps SrcSpan
l' (SrcSpan -> Name -> LHsExpr GhcRn
hsVar SrcSpan
l' Name
pureName) [LHsExpr GhcRn
x] where
    l' :: SrcSpan
l' = SrcSpan
noSrcSpan