content: -- 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 DeriveDataTypeable #-}
  {-# LANGUAGE DeriveFunctor #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE InstanceSigs #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeSynonymInstances #-}
  module Retrie.Universe
    ( Universe
    , printU
    , Matchable(..)
    , UMap(..)
    ) where
  
  import Control.Monad
  import Data.Data
  
  import Retrie.AlphaEnv
  import Retrie.ExactPrint
  import Retrie.GHC
  import Retrie.PatternMap.Class
  import Retrie.PatternMap.Instances
  import Retrie.Quantifiers
  import Retrie.Substitution
  
  -- | A sum type to collect all possible top-level rewritable types.
  data Universe
    = ULHsExpr (LHsExpr GhcPs)
    | ULStmt (LStmt GhcPs (LHsExpr GhcPs))
    | ULType (LHsType GhcPs)
    | ULPat (LPat GhcPs)
    deriving (Data)
  
  -- | Exactprint an annotated 'Universe'.
  printU :: Annotated Universe -> String
  printU u = exactPrintU (astA u)
      `debug` ("printU:" ++ showAst (astA u))
  
  -- | Primitive exactprint for 'Universe'.
  exactPrintU :: Universe -> String
  exactPrintU (ULHsExpr e) = exactPrint e
  exactPrintU (ULStmt s) = exactPrint s
  exactPrintU (ULType t) = exactPrint t
  exactPrintU (ULPat p) = exactPrint p
  
  -------------------------------------------------------------------------------
  
  -- | Class of types which can be injected into the 'Universe' type.
  class Matchable ast where
    -- | Inject an AST into 'Universe'
    inject :: ast -> Universe
  
    -- | Project an AST from a 'Universe'.
    -- Can fail if universe contains the wrong type.
    project :: Universe -> ast
  
    -- | Get the original location of the AST.
    getOrigin :: ast -> SrcSpan
  
  instance Matchable Universe where
    inject = id
    project = id
    getOrigin (ULHsExpr e) = getOrigin e
    getOrigin (ULStmt s) = getOrigin s
    getOrigin (ULType t) = getOrigin t
    getOrigin (ULPat p) = getOrigin p
  
  instance Matchable (LocatedA (HsExpr GhcPs)) where
    inject = ULHsExpr
    project (ULHsExpr x) = x
    project _ = error "project LHsExpr"
    getOrigin e = getLocA e
  
  instance Matchable (LocatedA (Stmt GhcPs (LocatedA (HsExpr GhcPs)))) where
    inject = ULStmt
    project (ULStmt x) = x
    project _ = error "project LStmt"
    getOrigin e = getLocA e
  
  instance Matchable (LocatedA (HsType GhcPs)) where
    inject = ULType
    project (ULType t) = t
    project _ = error "project ULType"
    getOrigin e = getLocA e
  
  instance Matchable (LocatedA (Pat GhcPs)) where
    inject = ULPat
    project (ULPat p) = p
    project _ = error "project ULPat"
    getOrigin = getLocA
  
  -------------------------------------------------------------------------------
  
  -- | The pattern map for 'Universe'.
  data UMap a = UMap
    { umExpr :: EMap a
    , umStmt :: SMap a
    , umType :: TyMap a
    , umPat  :: PatMap a
    }
    deriving (Functor)
  
  instance PatternMap UMap where
    type Key UMap = Universe
  
    mEmpty :: UMap a
    mEmpty = UMap mEmpty mEmpty mEmpty mEmpty
  
    mUnion :: UMap a -> UMap a -> UMap a
    mUnion m1 m2 = UMap
      (unionOn umExpr m1 m2)
      (unionOn umStmt m1 m2)
      (unionOn umType m1 m2)
      (unionOn umPat m1 m2)
  
    mAlter :: AlphaEnv -> Quantifiers -> Universe -> A a -> UMap a -> UMap a
    mAlter env vs u f m = go u
      where
        go (ULHsExpr e) = m { umExpr = mAlter env vs e f (umExpr m) }
        go (ULStmt s) = m { umStmt = mAlter env vs s f (umStmt m) }
        go (ULType t) = m { umType = mAlter env vs t f (umType m) }
        go (ULPat p) = m { umPat  = mAlter env vs (cLPat p) f (umPat m) }
  
    mMatch :: MatchEnv -> Universe -> (Substitution, UMap a) -> [(Substitution, a)]
    mMatch env = go
      where
        go (ULHsExpr e) = mapFor umExpr >=> mMatch env e
        go (ULStmt s) = mapFor umStmt >=> mMatch env s
        go (ULType t) = mapFor umType >=> mMatch env t
        go (ULPat p) = mapFor umPat >=> mMatch env (cLPat p)
  