-- Copyright (c) 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 RankNTypes #-}
module Retrie.Fixity
  ( FixityEnv
  , mkFixityEnv
  , lookupOp
  , lookupOpRdrName
  , Fixity(..)
  , FixityDirection(..)
  , extendFixityEnv
  , ppFixityEnv
  ) where

import Retrie.GHC

newtype FixityEnv = FixityEnv
  { FixityEnv -> FastStringEnv (FastString, Fixity)
unFixityEnv :: FastStringEnv (FastString, Fixity) }

instance Semigroup FixityEnv where
  -- | 'mappend' for 'FixityEnv' is right-biased
  (FixityEnv FastStringEnv (FastString, Fixity)
e1) <> :: FixityEnv -> FixityEnv -> FixityEnv
<> (FixityEnv FastStringEnv (FastString, Fixity)
e2) = FastStringEnv (FastString, Fixity) -> FixityEnv
FixityEnv (FastStringEnv (FastString, Fixity)
-> FastStringEnv (FastString, Fixity)
-> FastStringEnv (FastString, Fixity)
forall a. FastStringEnv a -> FastStringEnv a -> FastStringEnv a
plusFsEnv FastStringEnv (FastString, Fixity)
e1 FastStringEnv (FastString, Fixity)
e2)

instance Monoid FixityEnv where
  mempty :: FixityEnv
mempty = [(FastString, (FastString, Fixity))] -> FixityEnv
mkFixityEnv []

lookupOp :: LHsExpr GhcPs -> FixityEnv -> Fixity
lookupOp :: LHsExpr GhcPs -> FixityEnv -> Fixity
lookupOp (L SrcSpanAnnA
_ HsExpr GhcPs
e) | Just LIdP GhcPs
n <- HsExpr GhcPs -> Maybe (LIdP GhcPs)
forall p. HsExpr p -> Maybe (LIdP p)
varRdrName HsExpr GhcPs
e = LocatedN RdrName -> FixityEnv -> Fixity
lookupOpRdrName LIdP GhcPs
LocatedN RdrName
n
lookupOp LHsExpr GhcPs
_ = [Char] -> FixityEnv -> Fixity
forall a. HasCallStack => [Char] -> a
error [Char]
"lookupOp: called with non-variable!"

lookupOpRdrName :: LocatedN RdrName -> FixityEnv -> Fixity
lookupOpRdrName :: LocatedN RdrName -> FixityEnv -> Fixity
lookupOpRdrName (L SrcSpanAnnN
_ RdrName
n) (FixityEnv FastStringEnv (FastString, Fixity)
env) =
  Fixity
-> ((FastString, Fixity) -> Fixity)
-> Maybe (FastString, Fixity)
-> Fixity
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Fixity
defaultFixity (FastString, Fixity) -> Fixity
forall a b. (a, b) -> b
snd (Maybe (FastString, Fixity) -> Fixity)
-> Maybe (FastString, Fixity) -> Fixity
forall a b. (a -> b) -> a -> b
$ FastStringEnv (FastString, Fixity)
-> FastString -> Maybe (FastString, Fixity)
forall a. FastStringEnv a -> FastString -> Maybe a
lookupFsEnv FastStringEnv (FastString, Fixity)
env (OccName -> FastString
occNameFS (OccName -> FastString) -> OccName -> FastString
forall a b. (a -> b) -> a -> b
$ RdrName -> OccName
forall name. HasOccName name => name -> OccName
occName RdrName
n)

mkFixityEnv :: [(FastString, (FastString, Fixity))] -> FixityEnv
mkFixityEnv :: [(FastString, (FastString, Fixity))] -> FixityEnv
mkFixityEnv = FastStringEnv (FastString, Fixity) -> FixityEnv
FixityEnv (FastStringEnv (FastString, Fixity) -> FixityEnv)
-> ([(FastString, (FastString, Fixity))]
    -> FastStringEnv (FastString, Fixity))
-> [(FastString, (FastString, Fixity))]
-> FixityEnv
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(FastString, (FastString, Fixity))]
-> FastStringEnv (FastString, Fixity)
forall a. [(FastString, a)] -> FastStringEnv a
mkFsEnv

extendFixityEnv :: [(FastString, Fixity)] -> FixityEnv -> FixityEnv
extendFixityEnv :: [(FastString, Fixity)] -> FixityEnv -> FixityEnv
extendFixityEnv [(FastString, Fixity)]
l (FixityEnv FastStringEnv (FastString, Fixity)
env) =
  FastStringEnv (FastString, Fixity) -> FixityEnv
FixityEnv (FastStringEnv (FastString, Fixity) -> FixityEnv)
-> FastStringEnv (FastString, Fixity) -> FixityEnv
forall a b. (a -> b) -> a -> b
$ FastStringEnv (FastString, Fixity)
-> [(FastString, (FastString, Fixity))]
-> FastStringEnv (FastString, Fixity)
forall a. FastStringEnv a -> [(FastString, a)] -> FastStringEnv a
extendFsEnvList FastStringEnv (FastString, Fixity)
env [ (FastString
fs, (FastString, Fixity)
p) | p :: (FastString, Fixity)
p@(FastString
fs,Fixity
_) <- [(FastString, Fixity)]
l ]

ppFixityEnv :: FixityEnv -> String
ppFixityEnv :: FixityEnv -> [Char]
ppFixityEnv = [[Char]] -> [Char]
unlines ([[Char]] -> [Char])
-> (FixityEnv -> [[Char]]) -> FixityEnv -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((FastString, Fixity) -> [Char])
-> [(FastString, Fixity)] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (FastString, Fixity) -> [Char]
ppFixity ([(FastString, Fixity)] -> [[Char]])
-> (FixityEnv -> [(FastString, Fixity)]) -> FixityEnv -> [[Char]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FastStringEnv (FastString, Fixity) -> [(FastString, Fixity)]
forall key elt. UniqFM key elt -> [elt]
nonDetEltsUFM (FastStringEnv (FastString, Fixity) -> [(FastString, Fixity)])
-> (FixityEnv -> FastStringEnv (FastString, Fixity))
-> FixityEnv
-> [(FastString, Fixity)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FixityEnv -> FastStringEnv (FastString, Fixity)
unFixityEnv
  where
    ppFixity :: (FastString, Fixity) -> [Char]
ppFixity (FastString
fs, Fixity SourceText
_ Int
p FixityDirection
d) = [[Char]] -> [Char]
unwords
      [ case FixityDirection
d of
          FixityDirection
InfixN -> [Char]
"infix"
          FixityDirection
InfixL -> [Char]
"infixl"
          FixityDirection
InfixR -> [Char]
"infixr"
      , Int -> [Char]
forall a. Show a => a -> [Char]
show Int
p
      , FastString -> [Char]
unpackFS FastString
fs
      ]