    module Agda.TypeChecking.DiscrimTree.Types where

    import Control.DeepSeq

    import qualified Data.Map.Strict as Map
    import qualified Data.Set as Set

    import Data.Map.Strict (Map)
    import Data.Set (Set)

    import GHC.Generics (Generic)

    import Agda.Syntax.Internal
    import Agda.Syntax.Position
    import Agda.Syntax.Common

    import Agda.Utils.Impossible
    import Agda.Utils.Null

    data Key
      = RigidK {-# UNPACK #-} !QName {-# UNPACK #-} !Int
        -- ^ Rigid symbols (constructors, data types, record types,
        -- postulates) identified by a QName.
      | LocalK {-# UNPACK #-} !Int {-# UNPACK #-} !Int
        -- ^ Local variables.

      | PiK Hiding
        -- ^ Dependent function types. The domain will be represented
        -- accurately, for the case of a genuine dependent function type,
        -- the codomain will be a dummy.

      | ConstK
        -- ^ Constant lambdas.

      | SortK
        -- ^ Universes.

      | FlexK
      -- ^ Anything else.
      deriving (Int -> Key -> ShowS
    [Key] -> ShowS
    Key -> String
    (Int -> Key -> ShowS)
    -> (Key -> String) -> ([Key] -> ShowS) -> Show Key
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> Key -> ShowS
    showsPrec :: Int -> Key -> ShowS
    $cshow :: Key -> String
    show :: Key -> String
    $cshowList :: [Key] -> ShowS
    showList :: [Key] -> ShowS
    Show, Key -> Key -> Bool
    (Key -> Key -> Bool) -> (Key -> Key -> Bool) -> Eq Key
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: Key -> Key -> Bool
    == :: Key -> Key -> Bool
    $c/= :: Key -> Key -> Bool
    /= :: Key -> Key -> Bool
    Eq, Eq Key
    Eq Key =>
    (Key -> Key -> Ordering)
    -> (Key -> Key -> Bool)
    -> (Key -> Key -> Bool)
    -> (Key -> Key -> Bool)
    -> (Key -> Key -> Bool)
    -> (Key -> Key -> Key)
    -> (Key -> Key -> Key)
    -> Ord Key
    Key -> Key -> Bool
    Key -> Key -> Ordering
    Key -> Key -> Key
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: Key -> Key -> Ordering
    compare :: Key -> Key -> Ordering
    $c< :: Key -> Key -> Bool
    < :: Key -> Key -> Bool
    $c<= :: Key -> Key -> Bool
    <= :: Key -> Key -> Bool
    $c> :: Key -> Key -> Bool
    > :: Key -> Key -> Bool
    $c>= :: Key -> Key -> Bool
    >= :: Key -> Key -> Bool
    $cmax :: Key -> Key -> Key
    max :: Key -> Key -> Key
    $cmin :: Key -> Key -> Key
    min :: Key -> Key -> Key
    Ord, (forall x. Key -> Rep Key x)
    -> (forall x. Rep Key x -> Key) -> Generic Key
    forall x. Rep Key x -> Key
    forall x. Key -> Rep Key x
    forall a.
    (forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
    $cfrom :: forall x. Key -> Rep Key x
    from :: forall x. Key -> Rep Key x
    $cto :: forall x. Rep Key x -> Key
    to :: forall x. Rep Key x -> Key
    Generic)

    instance NFData Key

    -- | A 'Term'-indexed associative data structure supporting
    -- /approximate/ (conservative) lookup. Rather than using a @Trie@ keyed
    -- by 'Key' directly, a 'DiscrimTree' is instead represented more like a
    -- /case/ tree.
    --
    -- This allows us to exploit the fact that instance selection often
    -- focuses on a small part of the term: Only that critical chain is
    -- represented in the tree. As an example, level parameters are unlikely
    -- to contribute to narrowing a search problem, so it would be wasteful
    -- to have an indirection in the tree for every 'FlexK' standing for a
    -- level parameter.
    data DiscrimTree a
      = EmptyDT
        -- ^ The empty discrimination tree.
      | DoneDT (Set a)
        -- ^ Succeed with a given set of values.
      | CaseDT
        -- ^ Do case analysis on a term. 'CaseDT' is scoped in the same way
        -- as fast case trees for the abstract machine: When matching
        -- actually succeeds, the variable that was matched gets replaced by
        -- its arguments directly in the context.
          {-# UNPACK #-} !Int       -- ^ The variable to case on.
          (Map Key (DiscrimTree a)) -- ^ The proper branches.
          (DiscrimTree a)           -- ^ A further tree, which should always be explored.
      deriving ((forall x. DiscrimTree a -> Rep (DiscrimTree a) x)
    -> (forall x. Rep (DiscrimTree a) x -> DiscrimTree a)
    -> Generic (DiscrimTree a)
    forall x. Rep (DiscrimTree a) x -> DiscrimTree a
    forall x. DiscrimTree a -> Rep (DiscrimTree a) x
    forall a.
    (forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
    forall a x. Rep (DiscrimTree a) x -> DiscrimTree a
    forall a x. DiscrimTree a -> Rep (DiscrimTree a) x
    $cfrom :: forall a x. DiscrimTree a -> Rep (DiscrimTree a) x
    from :: forall x. DiscrimTree a -> Rep (DiscrimTree a) x
    $cto :: forall a x. Rep (DiscrimTree a) x -> DiscrimTree a
    to :: forall x. Rep (DiscrimTree a) x -> DiscrimTree a
    Generic, DiscrimTree a -> DiscrimTree a -> Bool
    (DiscrimTree a -> DiscrimTree a -> Bool)
    -> (DiscrimTree a -> DiscrimTree a -> Bool) -> Eq (DiscrimTree a)
    forall a. Eq a => DiscrimTree a -> DiscrimTree a -> Bool
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: forall a. Eq a => DiscrimTree a -> DiscrimTree a -> Bool
    == :: DiscrimTree a -> DiscrimTree a -> Bool
    $c/= :: forall a. Eq a => DiscrimTree a -> DiscrimTree a -> Bool
    /= :: DiscrimTree a -> DiscrimTree a -> Bool
    Eq, Int -> DiscrimTree a -> ShowS
    [DiscrimTree a] -> ShowS
    DiscrimTree a -> String
    (Int -> DiscrimTree a -> ShowS)
    -> (DiscrimTree a -> String)
    -> ([DiscrimTree a] -> ShowS)
    -> Show (DiscrimTree a)
    forall a. Show a => Int -> DiscrimTree a -> ShowS
    forall a. Show a => [DiscrimTree a] -> ShowS
    forall a. Show a => DiscrimTree a -> String
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: forall a. Show a => Int -> DiscrimTree a -> ShowS
    showsPrec :: Int -> DiscrimTree a -> ShowS
    $cshow :: forall a. Show a => DiscrimTree a -> String
    show :: DiscrimTree a -> String
    $cshowList :: forall a. Show a => [DiscrimTree a] -> ShowS
    showList :: [DiscrimTree a] -> ShowS
    Show)

    {-
    The extra continuation to CaseDT is used to represent instance tables
    which have non-trivial overlap, e.g.

      instance
        a : Foo X ?
        b : Foo ? X

    If we commited to the {a} branch of the discrimination tree, then we
    would miss {b} entirely. Note that an "obvious" overlap like

      instance
        a : Bar X
        b : Bar X

    would be represented as

      case 0 of
        Bar → case 0 of
          X → done {a, b}

    and the extra continuation would be empty.
    -}

    instance NFData a => NFData (DiscrimTree a)

    instance (KillRange a, Ord a) => KillRange (DiscrimTree a) where
      killRange :: KillRangeT (DiscrimTree a)
    killRange = \case
        DiscrimTree a
    EmptyDT      -> DiscrimTree a
    forall a. DiscrimTree a
    EmptyDT
        DoneDT Set a
    s     -> (Set a -> DiscrimTree a) -> Set a -> DiscrimTree a
    forall t (b :: Bool).
    (KILLRANGE t b, IsBase t ~ b, All KillRange (Domains t)) =>
    t -> t
    killRangeN Set a -> DiscrimTree a
    forall a. Set a -> DiscrimTree a
    DoneDT Set a
    s
        CaseDT Int
    i Map Key (DiscrimTree a)
    k DiscrimTree a
    o -> (Int -> Map Key (DiscrimTree a) -> KillRangeT (DiscrimTree a))
    -> Int -> Map Key (DiscrimTree a) -> KillRangeT (DiscrimTree a)
    forall t (b :: Bool).
    (KILLRANGE t b, IsBase t ~ b, All KillRange (Domains t)) =>
    t -> t
    killRangeN Int -> Map Key (DiscrimTree a) -> KillRangeT (DiscrimTree a)
    forall a.
    Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    CaseDT Int
    i Map Key (DiscrimTree a)
    k DiscrimTree a
    o

    instance Null (DiscrimTree a) where
      empty :: DiscrimTree a
    empty = DiscrimTree a
    forall a. DiscrimTree a
    EmptyDT
      null :: DiscrimTree a -> Bool
    null = \case
        DiscrimTree a
    EmptyDT -> Bool
    True
        DiscrimTree a
    _       -> Bool
    False

    -- | Merge a pair of discrimination trees. This function tries to build
    -- the minimal discrimination tree that yields the union of the inputs'
    -- results, though it does so slightly naïvely, without considerable
    -- optimisations (e.g. it does not turn single-alternative 'CaseDT's
    -- into 'DoneDT's).
    mergeDT :: Ord a => DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    mergeDT :: forall a. Ord a => DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    mergeDT DiscrimTree a
    EmptyDT    DiscrimTree a
    x = DiscrimTree a
    x
    mergeDT (DoneDT Set a
    s) DiscrimTree a
    x = case DiscrimTree a
    x of
      DiscrimTree a
    EmptyDT       -> Set a -> DiscrimTree a
    forall a. Set a -> DiscrimTree a
    DoneDT Set a
    s
      DoneDT Set a
    s'     -> Set a -> DiscrimTree a
    forall a. Set a -> DiscrimTree a
    DoneDT (Set a
    s Set a -> Set a -> Set a
    forall a. Semigroup a => a -> a -> a
    <> Set a
    s')
      CaseDT Int
    i Map Key (DiscrimTree a)
    bs DiscrimTree a
    x -> Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    forall a.
    Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    CaseDT Int
    i Map Key (DiscrimTree a)
    bs (DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    forall a. Ord a => DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    mergeDT (Set a -> DiscrimTree a
    forall a. Set a -> DiscrimTree a
    DoneDT Set a
    s) DiscrimTree a
    x)
    mergeDT (CaseDT Int
    i Map Key (DiscrimTree a)
    bs DiscrimTree a
    els) DiscrimTree a
    x = case DiscrimTree a
    x of
      DiscrimTree a
    EmptyDT  -> Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    forall a.
    Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    CaseDT Int
    i Map Key (DiscrimTree a)
    bs DiscrimTree a
    els
      DoneDT Set a
    s -> Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    forall a.
    Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    CaseDT Int
    i Map Key (DiscrimTree a)
    bs (DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    forall a. Ord a => DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    mergeDT (Set a -> DiscrimTree a
    forall a. Set a -> DiscrimTree a
    DoneDT Set a
    s) DiscrimTree a
    els)
      CaseDT Int
    j Map Key (DiscrimTree a)
    bs' DiscrimTree a
    els' -> case Int -> Int -> Ordering
    forall a. Ord a => a -> a -> Ordering
    compare Int
    i Int
    j of
        Ordering
    EQ -> Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    forall a.
    Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    CaseDT Int
    j ((DiscrimTree a -> DiscrimTree a -> DiscrimTree a)
    -> Map Key (DiscrimTree a)
    -> Map Key (DiscrimTree a)
    -> Map Key (DiscrimTree a)
    forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
    Map.unionWith DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    forall a. Ord a => DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    mergeDT Map Key (DiscrimTree a)
    bs Map Key (DiscrimTree a)
    bs') (DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    forall a. Ord a => DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    mergeDT DiscrimTree a
    els DiscrimTree a
    els')
        Ordering
    LT -> Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    forall a.
    Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    CaseDT Int
    i Map Key (DiscrimTree a)
    bs (DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    forall a. Ord a => DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    mergeDT DiscrimTree a
    els (Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    forall a.
    Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    CaseDT Int
    j Map Key (DiscrimTree a)
    bs' DiscrimTree a
    els'))
        Ordering
    GT -> Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    forall a.
    Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    CaseDT Int
    j Map Key (DiscrimTree a)
    bs' (DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    forall a. Ord a => DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    mergeDT DiscrimTree a
    els' (Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    forall a.
    Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    CaseDT Int
    i Map Key (DiscrimTree a)
    bs DiscrimTree a
    els))

    instance Ord a => Semigroup (DiscrimTree a) where
      <> :: DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    (<>) = DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    forall a. Ord a => DiscrimTree a -> DiscrimTree a -> DiscrimTree a
    mergeDT

    instance Ord a => Monoid (DiscrimTree a) where
      mempty :: DiscrimTree a
    mempty = DiscrimTree a
    forall a. DiscrimTree a
    EmptyDT

    -- | Construct the case tree corresponding to only performing proper
    -- matches on the given key. In this context, a "proper match" is any
    -- 'Key' that is not 'FlexK'.
    singletonDT :: [Key] -> a -> DiscrimTree a
    singletonDT :: forall a. [Key] -> a -> DiscrimTree a
    singletonDT [Key]
    key a
    val = Int -> [Key] -> DiscrimTree a
    go Int
    0 [Key]
    key where
      go :: Int -> [Key] -> DiscrimTree a
    go Int
    focus []         = Set a -> DiscrimTree a
    forall a. Set a -> DiscrimTree a
    DoneDT (Set a -> DiscrimTree a) -> Set a -> DiscrimTree a
    forall a b. (a -> b) -> a -> b
    $ a -> Set a
    forall a. a -> Set a
    Set.singleton a
    val
      go Int
    focus (Key
    FlexK:[Key]
    ts) = Int -> [Key] -> DiscrimTree a
    go (Int
    focus Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1) [Key]
    ts
      go Int
    focus (Key
    t:[Key]
    ts)     = Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    forall a.
    Int -> Map Key (DiscrimTree a) -> DiscrimTree a -> DiscrimTree a
    CaseDT Int
    focus (Key -> DiscrimTree a -> Map Key (DiscrimTree a)
    forall k a. k -> a -> Map k a
    Map.singleton Key
    t (Int -> [Key] -> DiscrimTree a
    go Int
    focus [Key]
    ts)) DiscrimTree a
    forall a. DiscrimTree a
    EmptyDT