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.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
    -- ^ 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