-- 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 (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 <- forall p. HsExpr p -> Maybe (LIdP p)
varRdrName HsExpr GhcPs
e = LocatedN RdrName -> FixityEnv -> Fixity
lookupOpRdrName LIdP GhcPs
n
lookupOp LHsExpr GhcPs
_ = 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) =
  forall b a. b -> (a -> b) -> Maybe a -> b
maybe Fixity
defaultFixity forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. FastStringEnv a -> FastString -> Maybe a
lookupFsEnv FastStringEnv (FastString, Fixity)
env (OccName -> FastString
occNameFS forall a b. (a -> b) -> a -> b
$ 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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 forall a b. (a -> b) -> a -> b
$ 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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (FastString, Fixity) -> [Char]
ppFixity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall key elt. UniqFM key elt -> [elt]
nonDetEltsUFM 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"
      , forall a. Show a => a -> [Char]
show Int
p
      , FastString -> [Char]
unpackFS FastString
fs
      ]