-- 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 CPP #-}
module Retrie.Substitution
  ( Substitution
  , HoleVal(..)
  , emptySubst
  , extendSubst
  , lookupSubst
  , deleteSubst
  , foldSubst
  ) where

import Retrie.ExactPrint
import Retrie.GHC

-- | A 'Substitution' is essentially a map from variable name to 'HoleVal'.
newtype Substitution = Substitution (UniqFM FastString (FastString, HoleVal))
-- See Note [Why not RdrNames?] for explanation of use of FastString

instance Show Substitution where
  show :: Substitution -> String
show (Substitution UniqFM FastString (FastString, HoleVal)
m) = [(FastString, HoleVal)] -> String
forall a. Show a => a -> String
show (UniqFM FastString (FastString, HoleVal) -> [(FastString, HoleVal)]
forall key elt. UniqFM key elt -> [elt]
nonDetEltsUFM UniqFM FastString (FastString, HoleVal)
m)

-- | Sum type of possible substitution values.
data HoleVal
  = HoleExpr AnnotatedHsExpr -- ^ 'HsExpr'
  | HolePat AnnotatedPat -- ^ 'Pat'
  | HoleType AnnotatedHsType -- ^ 'HsType'
  | HoleRdr RdrName -- ^ Alpha-renamed binder.

instance Show HoleVal where
  show :: HoleVal -> String
show (HoleExpr AnnotatedHsExpr
e) = String
"HoleExpr " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Annotated (GenLocated SrcSpanAnnA (HsExpr GhcPs)) -> String
forall ast. (Data ast, ExactPrint ast) => Annotated ast -> String
printA AnnotatedHsExpr
Annotated (GenLocated SrcSpanAnnA (HsExpr GhcPs))
e
  show (HolePat AnnotatedPat
p) = String
"HolePat " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Annotated (GenLocated SrcSpanAnnA (Pat GhcPs)) -> String
forall ast. (Data ast, ExactPrint ast) => Annotated ast -> String
printA AnnotatedPat
Annotated (GenLocated SrcSpanAnnA (Pat GhcPs))
p
  show (HoleType AnnotatedHsType
t) = String
"HoleType " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Annotated (GenLocated SrcSpanAnnA (HsType GhcPs)) -> String
forall ast. (Data ast, ExactPrint ast) => Annotated ast -> String
printA AnnotatedHsType
Annotated (GenLocated SrcSpanAnnA (HsType GhcPs))
t
  show (HoleRdr RdrName
r) = String
"HoleRdr " String -> ShowS
forall a. [a] -> [a] -> [a]
++ FastString -> String
unpackFS (RdrName -> FastString
rdrFS RdrName
r)

-- | The empty substitution.
emptySubst :: Substitution
emptySubst :: Substitution
emptySubst = UniqFM FastString (FastString, HoleVal) -> Substitution
Substitution UniqFM FastString (FastString, HoleVal)
forall key elt. UniqFM key elt
emptyUFM

-- | Lookup a value in the substitution.
lookupSubst :: FastString -> Substitution -> Maybe HoleVal
lookupSubst :: FastString -> Substitution -> Maybe HoleVal
lookupSubst FastString
k (Substitution UniqFM FastString (FastString, HoleVal)
m) = (FastString, HoleVal) -> HoleVal
forall a b. (a, b) -> b
snd ((FastString, HoleVal) -> HoleVal)
-> Maybe (FastString, HoleVal) -> Maybe HoleVal
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> UniqFM FastString (FastString, HoleVal)
-> FastString -> Maybe (FastString, HoleVal)
forall key elt. Uniquable key => UniqFM key elt -> key -> Maybe elt
lookupUFM UniqFM FastString (FastString, HoleVal)
m FastString
k

-- | Extend the substitution. If the key already exists, its value is replaced.
extendSubst :: Substitution -> FastString -> HoleVal -> Substitution
extendSubst :: Substitution -> FastString -> HoleVal -> Substitution
extendSubst (Substitution UniqFM FastString (FastString, HoleVal)
m) FastString
k HoleVal
v = UniqFM FastString (FastString, HoleVal) -> Substitution
Substitution (UniqFM FastString (FastString, HoleVal)
-> FastString
-> (FastString, HoleVal)
-> UniqFM FastString (FastString, HoleVal)
forall key elt.
Uniquable key =>
UniqFM key elt -> key -> elt -> UniqFM key elt
addToUFM UniqFM FastString (FastString, HoleVal)
m FastString
k (FastString
k,HoleVal
v))

-- | Delete from the substitution.
deleteSubst :: Substitution -> [FastString] -> Substitution
deleteSubst :: Substitution -> [FastString] -> Substitution
deleteSubst (Substitution UniqFM FastString (FastString, HoleVal)
m) [FastString]
ks = UniqFM FastString (FastString, HoleVal) -> Substitution
Substitution (UniqFM FastString (FastString, HoleVal)
-> [FastString] -> UniqFM FastString (FastString, HoleVal)
forall key elt.
Uniquable key =>
UniqFM key elt -> [key] -> UniqFM key elt
delListFromUFM UniqFM FastString (FastString, HoleVal)
m [FastString]
ks)

-- | Fold over the substitution.
foldSubst :: ((FastString, HoleVal) -> a -> a) -> a -> Substitution -> a
#if __GLASGOW_HASKELL__ < 908
foldSubst :: forall a.
((FastString, HoleVal) -> a -> a) -> a -> Substitution -> a
foldSubst (FastString, HoleVal) -> a -> a
f a
x (Substitution UniqFM FastString (FastString, HoleVal)
m) = ((FastString, HoleVal) -> a -> a)
-> a -> UniqFM FastString (FastString, HoleVal) -> a
forall elt a key. (elt -> a -> a) -> a -> UniqFM key elt -> a
foldUFM (FastString, HoleVal) -> a -> a
f a
x UniqFM FastString (FastString, HoleVal)
m
#else
foldSubst f x (Substitution m) = nonDetFoldUFM f x m
#endif