RewriteCase.hs

module Language.Haskell.Liquid.Constraint.RewriteCase 
    (getCaseRewrites) 
    where

import           Language.Fixpoint.Types
import qualified Language.Fixpoint.Misc as M
import           Language.Haskell.Liquid.Constraint.Types
import           Language.Haskell.Liquid.Types.Types
import           Language.Haskell.Liquid.Types.RType

import           Data.Maybe
import           Data.Tuple

import qualified Data.HashMap.Strict as M
import qualified Data.HashSet        as S

getCaseRewrites :: CGEnv -> SpecType -> LocalRewrites
getCaseRewrites γ spec =
    let Reft (_, refinement) = ur_reft $ rt_reft spec
        -- Names of constustors both type constructors and data constructors
        ctors   = toSet $ seBinds  $ constEnv γ
        -- All the global names (top level functions, types, etc...)
        globals = toSet $ reGlobal $ renv     γ
    in unloop
       $ concatMap (uncurry $ unify ctors globals)
       $ groupUnifiableEqualities 
       $ getEqualities refinement
    where toSet = S.fromList . M.keys

-- | Generates substitutions for non-global variables that make @e1@ and @e2@
-- equal.
--
-- If @v@ is not global, and @C@ is a data constructor
--
--  * @v@ and @e2@ produces @(v, e2)@
--  * @e1@ and @v@ produces @(e1, v)@
--  * @C a₁ ... aₙ@ and @C b₁ ... bₙ@ produces the substitutions from unifying
--    @(a₁, b₁), ..., (aₙ, bₙ)@
--
-- If any unification fails, the substitutions from the unifications that
-- succeed are still produced.
--
unify :: S.HashSet Symbol -> S.HashSet Symbol -> Expr -> Expr -> [(Symbol, Expr)]
unify ctors globals = go
    where
        go e1 e2 | e1 == e2 = []
        -- NOTE: We don't need to check for ECst because the expressions arent
        -- elaborated
        go (EVar s1) e2 | isLocal s1 = [(s1, e2)]
        go e1 (EVar s2) | isLocal s2 = [(s2, e1)]
        -- TODO: Tecnically we could also unify under lambdas but you have to be
        -- carefull about alpha equivalence idk if the effort is worth it.
        go e1 e2 
            -- Performing the unification under constructor is safe because 
            -- C a₁ ... aₙ = C b₁ ... bₙ ⟺ ∀ n . a₁ = bₙ
            | (EVar name1 , args1) <- splitEApp e2
            , (EVar name2 , args2) <- splitEApp e1
            , name1 == name2
            , isCtor name1
            , length args1 == length args2
            = concat $ zipWith go args1 args2
        go _ _ = []

        isCtor  name = name `S.member` ctors
        isLocal name = not (name `S.member` globals 
                           || name `S.member` ctors
                           || isPrefixOfSym anfPrefix name)


-- | Given a list of equalities this function produces the equalities that
-- result from applying transitivity exactly once. For instance, if we have
-- @[e1=e2, e2=e3, e1=e4]@ this function will produce @[e1=e3, e2=e4]@.
groupUnifiableEqualities :: [(Expr, Expr)] -> [(Expr, Expr)]
groupUnifiableEqualities = concat . concatMap mkEqs . grouping
  where
    mkEqs (e1 : es) = [ (e1, e) | e <- es ] : mkEqs es
    mkEqs _         = []

    grouping eqs = fmap snd $ M.groupList $ eqs ++ fmap swap eqs

getEqualities :: Expr -> [(Expr, Expr)]
getEqualities (PAtom Eq e1 e2) = [(e1, e2)]
getEqualities (PAnd eqs)       = concatMap getEqualities eqs
getEqualities _                = []


-- +-----------------------------------------------------+
-- | AcyclicRewrites: collection of rewrites that cannot |
-- | cause an infinite loop                              |
-- +-----------------------------------------------------+
-- This could be extracted as a separate module

-- | A collection of rewrites that cannot cause an infinite loop
newtype AcyclicRewrites = AR (M.HashMap Symbol Expr)

-- We can think of the map as a directed graph where every symbol is a vertex and
-- there is an edge v₁ → v₂ if v₂ is free in the expression that v₁ is rewritten to.
-- To guarantee that the set of rewrite rules is terminating, we ensure that there
-- aren't any cycles in the graph.

-- | Drops rewrites that would cause an infinite loop. The procedure is order
-- biased as rewrites earlier in the list take precedence.
unloop :: [(Symbol, Expr)] -> LocalRewrites
unloop = LocalRewrites . toRewrites . foldl' doInsert empty 
    where doInsert ar (s, e) = ar `fromMaybe` insert ar s e 

-- | Get the "raw" list of rewrites
toRewrites :: AcyclicRewrites -> M.HashMap Symbol Expr
toRewrites (AR m) = m

-- | @existsPth g s1 s2@ yields @True@ checks if there is a path from @s1@ to @s2@
-- in @g@. Time is @O(Σ(e ∈ m) |exprSymbolSet e|)@, or said otherwise, it is @O(m)@
-- where @m@ is the number of edges.
existsPath :: AcyclicRewrites -> Symbol -> Symbol -> Bool
existsPath (AR m) s1' s2 = go s1'
  where
    -- Since m is a DAG, we can use DFS to check if there is a path from s1 to
    -- s2 without worrying about infinite loops
    go s1 | s1 == s2 = True
    go s1 | Just e <- M.lookup s1 m
          = any go $ S.toList $ exprSymbolsSet e
    go _ = False

-- | Constructs an empty set of rewrites
empty :: AcyclicRewrites
empty = AR M.empty

-- | Inserts the rewrite if it wont't cause an infinite loop
insert :: AcyclicRewrites -> Symbol -> Expr -> Maybe AcyclicRewrites
insert ar s e | not $ s `S.member` exprSymbolsSet e
              , not $ any (\s2 -> existsPath ar s2 s) $ S.toList $ exprSymbolsSet e
              -- There are two ways to break the DAG invariant
              -- 1. If the rewrite is closing a loop
              -- 2. If the rewrite by itself is a cycle
              = Just $ insertUnsafe ar s e
              | otherwise 
              = Nothing
    where insertUnsafe (AR m) s' e' = AR $ M.insert s' e' m