content: {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  module Data.Abstract.Environment.Grammar
    ( Env
    , empty
    , fromList
    , fromList'
    , lookup
    , insert, insertList
    , insertRec, insertRecList
    , filter
    , union
    ) where
  
  import           Prelude hiding (lookup,filter,pred)
  
  import           Data.TreeGrammar (Grammar)
  import qualified Data.TreeGrammar as G
  import           Data.TreeGrammar.HashMap
  import           Data.TreeGrammar.Terminal (Terminal)
  import qualified Data.TreeGrammar.Terminal as T
  
  import           Data.Hashable
  import qualified Data.HashMap.Lazy as M
  import qualified Data.HashSet as H
  import           Data.Identifiable
  import           Data.Coerce
  import           Data.Empty
  
  newtype Env var val = Env (Grammar Int (HashMap var val))
    deriving (IsEmpty)
  
  instance (Identifiable var, Terminal val, Semigroup (val Int), Eq (val Int)) => Eq (Env var val) where
    Env env1 == Env env2 = G.identical env1 env2 || env1 == env2
  
  instance (Identifiable var, Terminal val, Eq (val Int), Semigroup (val Int)) => Hashable (Env var val) where
    hashWithSalt s (Env env) = G.hashTerminals s env
  
  deriving instance (Show var, Show (val Int)) => Show (Env var val)
  
  fromList :: (Identifiable var, Terminal val, Semigroup (val String), Semigroup (val Int)) => [(var, val ())] -> Env var val
  fromList l = Env $ G.grammar "Init" [("Init", HashMap (M.fromList [ (x,T.map (const "Init") v)| (x,v) <- l]))] []
  
  fromList' :: (Identifiable var, Terminal val) => [(var, val ())] -> Env var val -> Env var val
  fromList' = updateEnv (\start _ -> start) (\_ _ -> [])
  
  lookup :: (Identifiable var, Terminal val, Semigroup (val Int), Eq (val Int)) => var -> Env var val -> Maybe (val (Env var val))
  lookup var (Env env) =
    let HashMap m = G.toSubterms env
    in T.map Env <$> M.lookup var m
  
  insert :: (Identifiable var, Terminal val) => var -> (forall a. a -> val a)  -> Env var val -> Env var val
  insert var val = insertList [(var,val ())]
  
  insertList :: (Identifiable var, Terminal val) => [(var,val ())] -> Env var val -> Env var val
  insertList = updateEnv (\start _ -> start) (\start _ -> [start])
  
  insertRec :: (Identifiable var, Terminal val) => var -> (forall a. a -> val a) -> Env var val -> Env var val
  insertRec var val = insertRecList [(var,val ())]
  
  -- OPT: Cash largest non-terminal of the grammar to generate fresh non-terminals faster.
  insertRecList :: (Identifiable var, Terminal val) => [(var,val ())] -> Env var val -> Env var val
  insertRecList = updateEnv (\_ fresh -> fresh) (\start _ -> [start])
  
  union :: (Identifiable var, Terminal val, Semigroup (val Int)) => Env var val -> Env var val -> Env var val
  union (Env env1) (Env env2) = Env (G.union env1 env2)
  
  updateEnv :: (Terminal val, Identifiable var) => (Int -> Int -> Int) -> (Int -> Int -> [Int]) -> [(var, val ())] -> Env var val -> Env var val
  updateEnv _ _ [] env = env
  updateEnv closureEnv epsilon ls (Env env@G.Grammar{..}) = Env env
    { G.start = fresh
    , G.productions = M.insert fresh (G.Rhs { G.cons = coerce (M.fromList [ (var,T.map (const (closureEnv start fresh)) val)
                                                                          | (var,val) <- ls ])
                                            , G.eps = H.fromList (epsilon start fresh) }) productions
    }
    where
      fresh = succ $ M.foldlWithKey' (\x y _ -> max x y) 0 productions
  {-# INLINE updateEnv #-}
  
  filter :: (var -> Bool) -> Env var val -> Env var val
  filter pred (Env env@G.Grammar{..})= Env env
    { G.productions = M.map (\rhs -> rhs { G.cons = withHashMap (M.filterWithKey (\k _ -> pred k)) (G.cons rhs)}) productions
    }
  