content: {-# LANGUAGE CPP #-}
  module Subst(
               Subst,
               nullSubst, isNullSubst, (+->), mkSubst,
               Types(..), (@@), merge, mergeWith, mergeListWith,
               mergeAgreements,
               trimSubst, trimSubstByVars,
               {- removeFromSubst, -}
               apSubstToSubst,
               getSubstDomain, getSubstRange, sizeSubst,
               chkSubstOrder
               ) where
  
  import PPrint
  import CType
  import Type
  import Util(fromJustOrErr)
  import Data.List(nub,union)
  import Position
  
  import qualified Data.Set as Set
  import qualified Data.Map as Map
  import ErrorUtil(internalError)
  
  -- import Debug.Trace
  
  
  infixr 4 @@
  
  --we carry around a set of the variables on the right hand side
  --and a mapping from RHS variables to which LHS variable uses them
  
  type Set_TyVar = Set.Set TyVar
  type Set_vars = Map.Map TyVar Set_TyVar
  type S_map = Map.Map TyVar Type
  
  data Subst  = S S_map Set_vars
          deriving (Show, Eq)
  
  instance PPrint Subst where
      pPrint d p (S s v) = pparen (p>0) $ text "Subst" <+> pPrint d 0
                           ((Map.assocs s){-, printable_s_var v-} )
  
  
  nullSubst  :: Subst
  nullSubst   = S Map.empty Map.empty
  
  isNullSubst :: Subst -> Bool
  isNullSubst (S m _) = Map.null m
  
  (+->)      :: TyVar -> Type -> Subst
  u +-> t     = S (Map.singleton u t)
                  (back_set u t)
  
  back_set :: TyVar -> Type -> Set_vars
  back_set u t = (Map.fromList (map (\v -> (v, Set.singleton u)) (tv t)))
  
  mkSubst :: [(TyVar, Type)] -> Subst
  mkSubst vts = -- trace ("mkSubst: " ++ ppReadable vts) $
          S (Map.fromListWith avoid_duplicate_substitutions vts)
            (mk_set_vars vts)
  
  mk_set_vars :: [(TyVar, Type)] -> Set_vars
  mk_set_vars vts = (Map.unionsWith Set.union (map (uncurry back_set) vts))
  {- all calls to union (Map.union, Set.union, *.union*with) will be efficient
     because of the O(N+M) running time.  Repeated inserts would be less
     efficient because they are O(N*log(M)) time. -}
  
  -- check that concatenating a new substitution with an old one (via @@)
  -- will result in a valid substitution
  -- basically check that the RHS of the new substitution should not
  -- involve any variables in the LHS of the old substitution
  chkSubstOrder :: Subst -> Subst -> Bool
  chkSubstOrder (S _ new_rhs) (S old_lhs _) =
    Set.null (Set.intersection
               (Map.keysSet new_rhs)
               (Map.keysSet old_lhs))
  
  -- Add a newer substitution (ss1) to an older substitution (ss2).
  -- Any knowledge in ss1 is applied to ss2, and then the bindings are merged.
  --
  (@@) :: Subst -> Subst -> Subst
  ss1@(S s1 _) @@ ss2@(S _ var_old) =
      case (Set.null (Set.intersection (Map.keysSet var_old)
                                       (Map.keysSet s1))) of
       True -> --no need to apply s1 to s2
               --trace ("no ap: " ++ ppReadable (ss1, ss2)) $
                  merge_sv ss1 ss2
       _ -> -- trace ("ap: " ++ ppReadable (ss1, ss2, finished_2)) $
            (merge_sv (apSubstToSubst ss1 ss2) ss1)
  
  
  -- Apply the knowledge from one substitution to another, but without adding
  -- the bindings
  --
  apSubstToSubst :: Subst -> Subst -> Subst
  apSubstToSubst ss1@(S s1 _) ss2 = finished_2
            where
            finished_2 :: Subst
            finished_2 = Map.foldrWithKey fast_apply ss2 s1
            fast_apply ::  TyVar -> Type -> Subst -> Subst
            fast_apply v t a@(S ss vv) =
                case (Map.lookup v vv) of
                 Nothing -> a
                 Just lhs ->
  -- Forcing the apSub saves some memory (20%) at the cost of some time.
                  S (Set.fold (Map.adjust (seqCType . (apSub this_sub)))
                              ss lhs)
                    after_adds
                  where
                  this_sub = ( v +-> t)
                  delete_old :: Set_vars
                  --set left after variables that will get substituted are deleted
                  delete_old = Set.fold
                               (\l -> Map.update (possibly_delete l) v)
                               vv lhs
                  possibly_delete :: TyVar -> Set_TyVar -> Maybe Set_TyVar
                  possibly_delete l s = let
                                          r = Set.delete l s
                                        in case Set.null r of
                                            True -> Nothing
                                            _ -> Just r
                  after_adds :: Set_vars
                  -- set left after additions to the set delete_old
                  after_adds = foldr add_tyvar delete_old (tv t)
                  add_tyvar :: TyVar -> Set_vars  -> Set_vars
                  add_tyvar x set =
                        Map.unionWith Set.union
                           (Map.singleton x lhs) set
  
  avoid_duplicate_substitutions::Type -> Type -> Type
  avoid_duplicate_substitutions _ b = b
  --       case (a==b) of True -> a
  --                      _ ->error "duplicate different Tyvars in Subst"
  
  merge_s :: S_map -> S_map -> S_map
  merge_s = Map.unionWith avoid_duplicate_substitutions
  
  merge_v :: Set_vars -> Set_vars -> Set_vars
  merge_v = Map.unionWith Set.union
  
  merge_sv :: Subst -> Subst -> Subst
  merge_sv (S s1 var1) (S s2 var2) = (S (merge_s s1 s2)
                                      (merge_v var1 var2)
                                     )
  
  merge      :: Subst -> Subst -> Maybe Subst
  merge ss1 ss2 =
      if agree then return (merge_sv ss1 ss2)
        else fail "merge fails"
      where
         agree = and (map_on_the_intersection (==) ss1 ss2)
  
  map_on_the_intersection :: (Type -> Type -> a) -> Subst -> Subst -> [a]
  map_on_the_intersection call_f ss1@(S s1 _) ss2@(S s2 _) =
     map (\v -> call_f (apSub ss1 (TVar v))
                       (apSub ss2 (TVar v)))
         (Set.toList (Set.intersection (Map.keysSet s1)
                                       (Map.keysSet s2)))
  
  mergeListWith :: (Subst -> Subst -> Maybe Subst) -> [Maybe Subst] -> Maybe Subst
  mergeListWith merge_func = foldr cons (Just nullSubst)
   where cons :: Maybe Subst -> Maybe Subst -> Maybe Subst
         cons (Just s) (Just s') =
             -- rtrace ("mergeListWith: Just, Just: " ++ ppReadable (merge s s', s, s')) $
             merge_func s s'
         cons a        b         =
             -- rtrace ("mergeListWith: Nothing: " ++ ppReadable (a,b)) $
             Nothing
  
  -- the merge func is always unification
  -- it is a parameter because Unify imports subst
  -- this is why we do the @@ below - applying new knowledge to the subst
  -- is the right thing to do when unifying
  mergeWith :: (Type -> Type -> Maybe Subst) -> Subst -> Subst -> Maybe Subst
  mergeWith func ss1 ss2 =
    case merged_subts of
        Nothing -> Nothing
        -- need to apply anything we've learned to both substitutions
        Just ms -> Just (ms @@ (merge_sv ss1 ss2))
    where merged_subts::Maybe Subst
          merged_subts = mergeListWith (mergeWith func)
                             (map_on_the_intersection func ss1 ss2)
  
  -- merge substitutions, but drop mappings where there are disagreements
  -- used for typeclass defaulting
  mergeAgreements :: Subst -> Subst -> Subst
  mergeAgreements s1 s2 = fj $ merge diff1 diff2
    where agree s v t = v' == t || v' == TVar v
            where v' = apSub s (TVar v)
          diff1 = trimSubstBy (agree s2) s1
          diff2 = trimSubstBy (agree s1) s2
          fj = fromJustOrErr ("mergeAgreements: " ++
                              ppReadable (s1, s2, diff1, diff2))
  
  
  -------
  
  class Types t where
    apSub :: Subst -> t -> t
    tv    :: t -> [TyVar]
  
  instance Types Type where
    apSub (S seo _) v@(TVar u) =
          case slookup u seo of
          Just t  ->
              case t of
              (TGen _ _) -> t -- (kind (TGen _ _)) errors out -- don't check kind
              _ -> if isKVar (kind v) || kind t == kind v
                   then t
                   else internalError ("Subst.Type.apSub: bad kind: " ++
                                       ppString t ++ "::" ++ ppString (kind t) ++
                                       "@" ++ ppString (getPosition t) ++
                                       " / " ++ ppString v ++ "::" ++
                                       ppString (kind v) ++ "@" ++
                                       ppString (getPosition v))
          Nothing -> v
    apSub s (TAp l r) = normTAp (apSub s l) (apSub s r)
    apSub s t         = t
  
    tv (TVar u)  = [u]
    tv (TAp l r) = tv l `union` tv r
    tv t         = []
  
  instance Types a => Types [a] where
    apSub s ts = map (apSub s) ts
    tv ts      = nub (concatMap tv ts)
  
  slookup :: TyVar -> S_map -> Maybe Type
  slookup v xs = {-if length xs > 1300 then trace ("slookup " ++ unwords [itos x|(_,x,_)<-xs]) slookup' v xs else-} Map.lookup v xs
  
  -------
  
  -- These functions are useful for tracing and detecting internal errors
  -- on the substitution.
  
  sizeSubst :: Subst -> Int
  sizeSubst (S eo _) = (Map.size eo)
  
  getSubstDomain :: Subst -> [TyVar]
  getSubstDomain (S eo _) = (Map.keys eo)
  
  getSubstRange :: Subst -> [TyVar]
  getSubstRange (S eo _ ) = concat (map tv (Map.elems eo))
  
  
  -------
  
  -- This is a general substitution trimming function to unify
  -- the shared code from trimSubst and trimSubstByVars
  -- it also uses map and set operations to compute the new
  -- substitution efficiently
  
  fixupBack :: S_map -> Set_vars -> Set_vars
  fixupBack new_map old_back = Map.filter
                                 (not . Set.null)
                                 (Map.map (Set.intersection retained_vars) old_back)
    where retained_vars = Map.keysSet new_map
  
  -- It takes a function that determines what elements of the
  -- substitution to keep and fixes the subst and the back set
  trimSubstBy :: (TyVar -> Type -> Bool) -> Subst -> Subst
  trimSubstBy filterFunc (S old_map old_back) = S new_map new_back
    where new_map       = Map.filterWithKey filterFunc old_map
          new_back      = fixupBack new_map old_back
  
  -- This trims substitutions for all variables above a certain number.
  -- It is used in type-checking explicitly typed let-bindings, in order
  -- to roll the substitution back and throw away all the new associations
  -- introduced for the let-binding (which are not needed again).
  -- It takes advantage of the fact that new type variables are created
  -- in numeric order.
  trimSubst :: TyVar -> Subst -> Subst
  trimSubst v (S old_map old_back) = S new_map new_back
    where (new_map, _) = Map.split v old_map
          new_back     = fixupBack new_map old_back
  
  -- This removes substitution entries whose LHS or RHS contains the variables.
  -- (That is, entried for a type variable and whose associated type contains
  -- the variable.)
  -- It is used to removed mention of type variables inside a type lambda
  -- (to prevent name capture on the lambda-bound variables when substituting
  -- inside the lambda).
  
  trimSubstByVars :: [TyVar] -> Subst -> Subst
  trimSubstByVars vs = if (null vs) then id else trimSubstBy filterFunc
    where filterFunc v t = not ((elem v vs) ||
                                (any (\u -> elem u vs) (tv t)))
  
  {-
  -- This removes variables just from the LHS of the substitution,
  -- and returns a list of the variables which were found.
  removeFromSubst :: [TyVar] -> Subst -> (Subst, [TyVar])
  removeFromSubst vs (S e o) =
      let
         -- removes and returns a list of the removed vars
         removeFunc s@(v,_,_) (ss, removed_vs) =
             if (elem v vs)
             then (ss, v:removed_vs)
             else (s:ss, removed_vs)
  
         (e',rvs1) = foldr removeFunc ([],[]) e
         (o',rvs2) = foldr removeFunc ([],rvs1) o
      in
         (S e' o', rvs2)
  -}
  