{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Ormolu.Printer.Operators
( OpTree (..),
opTreeLoc,
reassociateOpTree,
)
where
import Data.Function (on)
import qualified Data.List as L
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import Data.Maybe (fromMaybe, mapMaybe)
import GHC.Types.Basic
import GHC.Types.Name.Occurrence (occNameString)
import GHC.Types.Name.Reader
import GHC.Types.SrcLoc
import Ormolu.Utils (unSrcSpan)
data OpTree ty op
= OpNode ty
| OpBranch
(OpTree ty op)
op
(OpTree ty op)
opTreeLoc :: OpTree (Located a) b -> SrcSpan
opTreeLoc :: OpTree (Located a) b -> SrcSpan
opTreeLoc (OpNode (L SrcSpan
l a
_)) = SrcSpan
l
opTreeLoc (OpBranch OpTree (Located a) b
l b
_ OpTree (Located a) b
r) = SrcSpan -> SrcSpan -> SrcSpan
combineSrcSpans (OpTree (Located a) b -> SrcSpan
forall a b. OpTree (Located a) b -> SrcSpan
opTreeLoc OpTree (Located a) b
l) (OpTree (Located a) b -> SrcSpan
forall a b. OpTree (Located a) b -> SrcSpan
opTreeLoc OpTree (Located a) b
r)
reassociateOpTree ::
(op -> Maybe RdrName) ->
OpTree (Located ty) (Located op) ->
OpTree (Located ty) (Located op)
reassociateOpTree :: (op -> Maybe RdrName)
-> OpTree (Located ty) (Located op)
-> OpTree (Located ty) (Located op)
reassociateOpTree op -> Maybe RdrName
getOpName OpTree (Located ty) (Located op)
opTree =
Map String Fixity
-> (Located op -> Maybe RdrName)
-> OpTree (Located ty) (Located op)
-> OpTree (Located ty) (Located op)
forall ty op.
Map String Fixity
-> (op -> Maybe RdrName) -> OpTree ty op -> OpTree ty op
reassociateOpTreeWith
((op -> Maybe RdrName)
-> OpTree (Located ty) (Located op) -> Map String Fixity
forall ty op.
(op -> Maybe RdrName)
-> OpTree (Located ty) (Located op) -> Map String Fixity
buildFixityMap op -> Maybe RdrName
getOpName OpTree (Located ty) (Located op)
normOpTree)
(op -> Maybe RdrName
getOpName (op -> Maybe RdrName)
-> (Located op -> op) -> Located op -> Maybe RdrName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Located op -> op
forall l e. GenLocated l e -> e
unLoc)
OpTree (Located ty) (Located op)
normOpTree
where
normOpTree :: OpTree (Located ty) (Located op)
normOpTree = OpTree (Located ty) (Located op)
-> OpTree (Located ty) (Located op)
forall ty op. OpTree ty op -> OpTree ty op
normalizeOpTree OpTree (Located ty) (Located op)
opTree
reassociateOpTreeWith ::
forall ty op.
Map String Fixity ->
(op -> Maybe RdrName) ->
OpTree ty op ->
OpTree ty op
reassociateOpTreeWith :: Map String Fixity
-> (op -> Maybe RdrName) -> OpTree ty op -> OpTree ty op
reassociateOpTreeWith Map String Fixity
fixityMap op -> Maybe RdrName
getOpName = OpTree ty op -> OpTree ty op
go
where
fixityOf :: op -> Fixity
fixityOf :: op -> Fixity
fixityOf op
op = Fixity -> Maybe Fixity -> Fixity
forall a. a -> Maybe a -> a
fromMaybe Fixity
defaultFixity (Maybe Fixity -> Fixity) -> Maybe Fixity -> Fixity
forall a b. (a -> b) -> a -> b
$ do
String
s <- OccName -> String
occNameString (OccName -> String) -> (RdrName -> OccName) -> RdrName -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RdrName -> OccName
rdrNameOcc (RdrName -> String) -> Maybe RdrName -> Maybe String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> op -> Maybe RdrName
getOpName op
op
String -> Map String Fixity -> Maybe Fixity
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup String
s Map String Fixity
fixityMap
go :: OpTree ty op -> OpTree ty op
go :: OpTree ty op -> OpTree ty op
go t :: OpTree ty op
t@(OpNode ty
_) = OpTree ty op
t
go t :: OpTree ty op
t@(OpBranch (OpNode ty
_) op
_ (OpNode ty
_)) = OpTree ty op
t
go (OpBranch l :: OpTree ty op
l@(OpNode ty
_) op
op (OpBranch OpTree ty op
l' op
op' OpTree ty op
r')) =
OpTree ty op -> OpTree ty op
go (OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch (OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch OpTree ty op
l op
op OpTree ty op
l') op
op' OpTree ty op
r')
go (OpBranch (OpBranch OpTree ty op
l op
op OpTree ty op
r) op
op' r' :: OpTree ty op
r'@(OpNode ty
_)) =
if (Bool, Bool) -> Bool
forall a b. (a, b) -> b
snd ((Bool, Bool) -> Bool) -> (Bool, Bool) -> Bool
forall a b. (a -> b) -> a -> b
$ Fixity -> Fixity -> (Bool, Bool)
compareFixity (op -> Fixity
fixityOf op
op) (op -> Fixity
fixityOf op
op')
then OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch OpTree ty op
l op
op (OpTree ty op -> OpTree ty op
go (OpTree ty op -> OpTree ty op) -> OpTree ty op -> OpTree ty op
forall a b. (a -> b) -> a -> b
$ OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch OpTree ty op
r op
op' OpTree ty op
r')
else OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch (OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch OpTree ty op
l op
op OpTree ty op
r) op
op' OpTree ty op
r'
go (OpBranch (OpBranch OpTree ty op
l op
op OpTree ty op
r) op
op' (OpBranch OpTree ty op
l' op
op'' OpTree ty op
r')) =
if (Bool, Bool) -> Bool
forall a b. (a, b) -> b
snd ((Bool, Bool) -> Bool) -> (Bool, Bool) -> Bool
forall a b. (a -> b) -> a -> b
$ Fixity -> Fixity -> (Bool, Bool)
compareFixity (op -> Fixity
fixityOf op
op) (op -> Fixity
fixityOf op
op')
then OpTree ty op -> OpTree ty op
go (OpTree ty op -> OpTree ty op) -> OpTree ty op -> OpTree ty op
forall a b. (a -> b) -> a -> b
$ OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch (OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch OpTree ty op
l op
op (OpTree ty op -> OpTree ty op
go (OpTree ty op -> OpTree ty op) -> OpTree ty op -> OpTree ty op
forall a b. (a -> b) -> a -> b
$ OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch OpTree ty op
r op
op' OpTree ty op
l')) op
op'' OpTree ty op
r'
else OpTree ty op -> OpTree ty op
go (OpTree ty op -> OpTree ty op) -> OpTree ty op -> OpTree ty op
forall a b. (a -> b) -> a -> b
$ OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch (OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch (OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch OpTree ty op
l op
op OpTree ty op
r) op
op' OpTree ty op
l') op
op'' OpTree ty op
r'
data Score
=
AtBeginning Int
|
AtEnd
|
InBetween
deriving (Score -> Score -> Bool
(Score -> Score -> Bool) -> (Score -> Score -> Bool) -> Eq Score
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Score -> Score -> Bool
$c/= :: Score -> Score -> Bool
== :: Score -> Score -> Bool
$c== :: Score -> Score -> Bool
Eq, Eq Score
Eq Score
-> (Score -> Score -> Ordering)
-> (Score -> Score -> Bool)
-> (Score -> Score -> Bool)
-> (Score -> Score -> Bool)
-> (Score -> Score -> Bool)
-> (Score -> Score -> Score)
-> (Score -> Score -> Score)
-> Ord Score
Score -> Score -> Bool
Score -> Score -> Ordering
Score -> Score -> Score
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Score -> Score -> Score
$cmin :: Score -> Score -> Score
max :: Score -> Score -> Score
$cmax :: Score -> Score -> Score
>= :: Score -> Score -> Bool
$c>= :: Score -> Score -> Bool
> :: Score -> Score -> Bool
$c> :: Score -> Score -> Bool
<= :: Score -> Score -> Bool
$c<= :: Score -> Score -> Bool
< :: Score -> Score -> Bool
$c< :: Score -> Score -> Bool
compare :: Score -> Score -> Ordering
$ccompare :: Score -> Score -> Ordering
$cp1Ord :: Eq Score
Ord)
buildFixityMap ::
forall ty op.
(op -> Maybe RdrName) ->
OpTree (Located ty) (Located op) ->
Map String Fixity
buildFixityMap :: (op -> Maybe RdrName)
-> OpTree (Located ty) (Located op) -> Map String Fixity
buildFixityMap op -> Maybe RdrName
getOpName OpTree (Located ty) (Located op)
opTree =
Map String Fixity -> Map String Fixity
addOverrides
(Map String Fixity -> Map String Fixity)
-> ([(String, Score)] -> Map String Fixity)
-> [(String, Score)]
-> Map String Fixity
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(String, Fixity)] -> Map String Fixity
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
([(String, Fixity)] -> Map String Fixity)
-> ([(String, Score)] -> [(String, Fixity)])
-> [(String, Score)]
-> Map String Fixity
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, [(String, Score)]) -> [(String, Fixity)])
-> [(Int, [(String, Score)])] -> [(String, Fixity)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(Int
i, [(String, Score)]
ns) -> ((String, Score) -> (String, Fixity))
-> [(String, Score)] -> [(String, Fixity)]
forall a b. (a -> b) -> [a] -> [b]
map (\(String
n, Score
_) -> (String
n, Int -> FixityDirection -> Fixity
fixity Int
i FixityDirection
InfixL)) [(String, Score)]
ns)
([(Int, [(String, Score)])] -> [(String, Fixity)])
-> ([(String, Score)] -> [(Int, [(String, Score)])])
-> [(String, Score)]
-> [(String, Fixity)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [[(String, Score)]] -> [(Int, [(String, Score)])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
2 ..]
([[(String, Score)]] -> [(Int, [(String, Score)])])
-> ([(String, Score)] -> [[(String, Score)]])
-> [(String, Score)]
-> [(Int, [(String, Score)])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((String, Score) -> (String, Score) -> Bool)
-> [(String, Score)] -> [[(String, Score)]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
L.groupBy (Score -> Score -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Score -> Score -> Bool)
-> ((String, Score) -> Score)
-> (String, Score)
-> (String, Score)
-> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (String, Score) -> Score
forall a b. (a, b) -> b
snd)
([(String, Score)] -> [[(String, Score)]])
-> ([(String, Score)] -> [(String, Score)])
-> [(String, Score)]
-> [[(String, Score)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(String, Score)] -> [(String, Score)]
selectScores
([(String, Score)] -> Map String Fixity)
-> [(String, Score)] -> Map String Fixity
forall a b. (a -> b) -> a -> b
$ OpTree (Located ty) (Located op) -> [(String, Score)]
score OpTree (Located ty) (Located op)
opTree
where
addOverrides :: Map String Fixity -> Map String Fixity
addOverrides :: Map String Fixity -> Map String Fixity
addOverrides Map String Fixity
m =
[(String, Fixity)] -> Map String Fixity
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
[ (String
"$", Int -> FixityDirection -> Fixity
fixity Int
0 FixityDirection
InfixR),
(String
":", Int -> FixityDirection -> Fixity
fixity Int
1 FixityDirection
InfixR),
(String
".", Int -> FixityDirection -> Fixity
fixity Int
100 FixityDirection
InfixL)
]
Map String Fixity -> Map String Fixity -> Map String Fixity
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Map String Fixity
m
fixity :: Int -> FixityDirection -> Fixity
fixity = SourceText -> Int -> FixityDirection -> Fixity
Fixity SourceText
NoSourceText
score :: OpTree (Located ty) (Located op) -> [(String, Score)]
score :: OpTree (Located ty) (Located op) -> [(String, Score)]
score (OpNode Located ty
_) = []
score (OpBranch OpTree (Located ty) (Located op)
l Located op
o OpTree (Located ty) (Located op)
r) = [(String, Score)] -> Maybe [(String, Score)] -> [(String, Score)]
forall a. a -> Maybe a -> a
fromMaybe (OpTree (Located ty) (Located op) -> [(String, Score)]
score OpTree (Located ty) (Located op)
r) (Maybe [(String, Score)] -> [(String, Score)])
-> Maybe [(String, Score)] -> [(String, Score)]
forall a b. (a -> b) -> a -> b
$ do
Int
le <- RealSrcSpan -> Int
srcSpanEndLine (RealSrcSpan -> Int) -> Maybe RealSrcSpan -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcSpan -> Maybe RealSrcSpan
unSrcSpan (OpTree (Located ty) (Located op) -> SrcSpan
forall a b. OpTree (Located a) b -> SrcSpan
opTreeLoc OpTree (Located ty) (Located op)
l)
Int
ob <- RealSrcSpan -> Int
srcSpanStartLine (RealSrcSpan -> Int) -> Maybe RealSrcSpan -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcSpan -> Maybe RealSrcSpan
unSrcSpan (Located op -> SrcSpan
forall l e. GenLocated l e -> l
getLoc Located op
o)
Int
oe <- RealSrcSpan -> Int
srcSpanEndLine (RealSrcSpan -> Int) -> Maybe RealSrcSpan -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcSpan -> Maybe RealSrcSpan
unSrcSpan (Located op -> SrcSpan
forall l e. GenLocated l e -> l
getLoc Located op
o)
Int
rb <- RealSrcSpan -> Int
srcSpanStartLine (RealSrcSpan -> Int) -> Maybe RealSrcSpan -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcSpan -> Maybe RealSrcSpan
unSrcSpan (OpTree (Located ty) (Located op) -> SrcSpan
forall a b. OpTree (Located a) b -> SrcSpan
opTreeLoc OpTree (Located ty) (Located op)
r)
Int
oc <- RealSrcSpan -> Int
srcSpanStartCol (RealSrcSpan -> Int) -> Maybe RealSrcSpan -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcSpan -> Maybe RealSrcSpan
unSrcSpan (Located op -> SrcSpan
forall l e. GenLocated l e -> l
getLoc Located op
o)
String
opName <- OccName -> String
occNameString (OccName -> String) -> (RdrName -> OccName) -> RdrName -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RdrName -> OccName
rdrNameOcc (RdrName -> String) -> Maybe RdrName -> Maybe String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> op -> Maybe RdrName
getOpName (Located op -> op
forall l e. GenLocated l e -> e
unLoc Located op
o)
let s :: Score
s
| Int
le Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ob = Int -> Score
AtBeginning Int
oc
| Int
oe Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
rb = Score
AtEnd
| Bool
otherwise = Score
InBetween
[(String, Score)] -> Maybe [(String, Score)]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(String, Score)] -> Maybe [(String, Score)])
-> [(String, Score)] -> Maybe [(String, Score)]
forall a b. (a -> b) -> a -> b
$ (String
opName, Score
s) (String, Score) -> [(String, Score)] -> [(String, Score)]
forall a. a -> [a] -> [a]
: OpTree (Located ty) (Located op) -> [(String, Score)]
score OpTree (Located ty) (Located op)
r
selectScores :: [(String, Score)] -> [(String, Score)]
selectScores :: [(String, Score)] -> [(String, Score)]
selectScores =
((String, Score) -> Score)
-> [(String, Score)] -> [(String, Score)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
L.sortOn (String, Score) -> Score
forall a b. (a, b) -> b
snd
([(String, Score)] -> [(String, Score)])
-> ([(String, Score)] -> [(String, Score)])
-> [(String, Score)]
-> [(String, Score)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(String, Score)] -> Maybe (String, Score))
-> [[(String, Score)]] -> [(String, Score)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
( \case
[] -> Maybe (String, Score)
forall a. Maybe a
Nothing
xs :: [(String, Score)]
xs@((String
n, Score
_) : [(String, Score)]
_) -> (String, Score) -> Maybe (String, Score)
forall a. a -> Maybe a
Just (String
n, [Score] -> Score
selectScore ([Score] -> Score) -> [Score] -> Score
forall a b. (a -> b) -> a -> b
$ ((String, Score) -> Score) -> [(String, Score)] -> [Score]
forall a b. (a -> b) -> [a] -> [b]
map (String, Score) -> Score
forall a b. (a, b) -> b
snd [(String, Score)]
xs)
)
([[(String, Score)]] -> [(String, Score)])
-> ([(String, Score)] -> [[(String, Score)]])
-> [(String, Score)]
-> [(String, Score)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((String, Score) -> (String, Score) -> Bool)
-> [(String, Score)] -> [[(String, Score)]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
L.groupBy (String -> String -> Bool
forall a. Eq a => a -> a -> Bool
(==) (String -> String -> Bool)
-> ((String, Score) -> String)
-> (String, Score)
-> (String, Score)
-> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (String, Score) -> String
forall a b. (a, b) -> a
fst)
([(String, Score)] -> [[(String, Score)]])
-> ([(String, Score)] -> [(String, Score)])
-> [(String, Score)]
-> [[(String, Score)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(String, Score)] -> [(String, Score)]
forall a. Ord a => [a] -> [a]
L.sort
selectScore :: [Score] -> Score
selectScore :: [Score] -> Score
selectScore [Score]
xs =
case (Score -> Bool) -> [Score] -> [Score]
forall a. (a -> Bool) -> [a] -> [a]
filter (Score -> Score -> Bool
forall a. Eq a => a -> a -> Bool
/= Score
InBetween) [Score]
xs of
[] -> Score
InBetween
[Score]
xs' -> [Score] -> Score
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Score]
xs'
normalizeOpTree :: OpTree ty op -> OpTree ty op
normalizeOpTree :: OpTree ty op -> OpTree ty op
normalizeOpTree (OpNode ty
n) =
ty -> OpTree ty op
forall ty op. ty -> OpTree ty op
OpNode ty
n
normalizeOpTree (OpBranch (OpNode ty
l) op
lop OpTree ty op
r) =
OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch (ty -> OpTree ty op
forall ty op. ty -> OpTree ty op
OpNode ty
l) op
lop (OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> OpTree ty op
normalizeOpTree OpTree ty op
r)
normalizeOpTree (OpBranch (OpBranch OpTree ty op
l' op
lop' OpTree ty op
r') op
lop OpTree ty op
r) =
OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> OpTree ty op
normalizeOpTree (OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch OpTree ty op
l' op
lop' (OpTree ty op -> op -> OpTree ty op -> OpTree ty op
forall ty op. OpTree ty op -> op -> OpTree ty op -> OpTree ty op
OpBranch OpTree ty op
r' op
lop OpTree ty op
r))