-- | A term index to accelerate matching.
-- An index is a multimap from terms to arbitrary values.
--
-- The type of query supported is: given a search term, find all keys such that
-- the search term is an instance of the key, and return the corresponding
-- values.

{-# LANGUAGE BangPatterns, RecordWildCards, OverloadedStrings, FlexibleContexts, CPP, TupleSections, TypeFamilies #-}
-- We get some bogus warnings because of pattern synonyms.
{-# OPTIONS_GHC -fno-warn-overlapping-patterns #-}
{-# OPTIONS_GHC -O2 -fmax-worker-args=100 #-}
{-# OPTIONS_GHC -funfolding-use-threshold=1000 #-}
#ifdef USE_LLVM
{-# OPTIONS_GHC -fllvm #-}
#endif
module Twee.Index(
  Index,
  empty,
  null,
  singleton,
  insert,
  delete,
  lookup,
  matches,
  elems,
  fromList,
  fromListWith,
  invariant) where

import Prelude hiding (null, lookup)
import Twee.Base hiding (var, fun, empty, singleton, prefix, funs, lookupList, lookup, at)
import qualified Twee.Term as Term
import Data.DynamicArray hiding (singleton)
import qualified Data.DynamicArray as Array
import qualified Data.List as List
import Data.Numbered(Numbered)
import qualified Data.Numbered as Numbered
import qualified Data.IntMap.Strict as IntMap
import qualified Twee.Term.Core as Core
import Twee.Profile

-- The term index in this module is a _perfect discrimination tree_.
-- This is a trie whose keys are terms, represented as flat lists of symbols
-- (either functions or variables).
--
-- To insert a key-value pair into the discrimination tree, we do
-- ordinary trie insertion, but first canonicalising the key-value
-- pair so that variables are introduced in ascending order.
-- This canonicalisation reduces the size of the trie, but is also
-- needed for our particular implementation of lookup to work correctly
-- (specifically the extendBindings function below).
--
-- Lookup maintains a term list, which is initially the search term,
-- and a substitution. It proceeds down the trie, consuming bits of
-- the term list as it goes.
--
-- If the current trie node has an edge for a function symbol f, and the term at
-- the head of the term list is f(t1..tn), we can follow the f edge. We then
-- delete f from the term list, but keep t1..tn at the front of the term list.
-- (In other words we delete only the symbol f and not its arguments.)
--
-- If the current trie node has a variable edge x, we can follow that edge.
-- We remove the head term from the term list, as 'x' matches that
-- whole term. We then add the binding x->t to the substitution.
-- If the substitution already has a binding x->u with u/=t, we can't
-- follow the edge.
--
-- If the term list ever becomes empty, we have a match, and return
-- the substitution.
--
-- Often there are several edges we can follow (function symbol and
-- any number of variable edges), and in that case the algorithm uses
-- backtracking.

----------------------------------------------------------------------
-- The term index.
----------------------------------------------------------------------

-- | A term index: a multimap from @'Term' f@ to @a@.
data Index f a =
  -- A non-empty index.
  Index {
    -- The size of the smallest term in the index.
    forall f a. Index f a -> Int
minSize_ :: {-# UNPACK #-} !Int,
    -- When all keys in the index start with the same sequence of symbols, we
    -- compress them into this prefix; the "fun" and "var" fields below refer to
    -- the first symbol _after_ the prefix, and the "here" field contains values
    -- whose remaining key is exactly this prefix.
    forall f a. Index f a -> TermList f
prefix   :: {-# UNPACK #-} !(TermList f),
    -- The values that are found at this node.
    forall f a. Index f a -> [a]
here     :: [a],
    -- Function symbol edges.
    -- The array is indexed by function number.
    forall f a. Index f a -> Array (Index f a)
fun      :: {-# UNPACK #-} !(Array (Index f a)),
    -- List of variable edges indexed by variable number.
    -- Invariant: all edges present in the list are non-Nil.
    --
    -- Invariant: variables in terms are introduced in ascending
    -- order, with no gaps (i.e. if the term so far has the variables
    -- x1..xn, then the edges here must be drawn from x1...x{n+1}).
    forall f a. Index f a -> Numbered (Index f a)
var      :: {-# UNPACK #-} !(Numbered (Index f a)) } |
  -- An empty index.
  Nil
  deriving Int -> Index f a -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall f a.
(Labelled f, Show f, Show a) =>
Int -> Index f a -> ShowS
forall f a. (Labelled f, Show f, Show a) => [Index f a] -> ShowS
forall f a. (Labelled f, Show f, Show a) => Index f a -> String
showList :: [Index f a] -> ShowS
$cshowList :: forall f a. (Labelled f, Show f, Show a) => [Index f a] -> ShowS
show :: Index f a -> String
$cshow :: forall f a. (Labelled f, Show f, Show a) => Index f a -> String
showsPrec :: Int -> Index f a -> ShowS
$cshowsPrec :: forall f a.
(Labelled f, Show f, Show a) =>
Int -> Index f a -> ShowS
Show

minSize :: Index f a -> Int
minSize :: forall f a. Index f a -> Int
minSize Index f a
Nil = forall a. Bounded a => a
maxBound
minSize Index f a
idx = forall f a. Index f a -> Int
minSize_ Index f a
idx

-- | Check the invariant of an index. For debugging purposes.
invariant :: Index f a -> Bool
invariant :: forall f a. Index f a -> Bool
invariant Index f a
Nil = Bool
True
invariant Index{Int
[a]
Array (Index f a)
Numbered (Index f a)
TermList f
var :: Numbered (Index f a)
fun :: Array (Index f a)
here :: [a]
prefix :: TermList f
minSize_ :: Int
var :: forall f a. Index f a -> Numbered (Index f a)
fun :: forall f a. Index f a -> Array (Index f a)
here :: forall f a. Index f a -> [a]
prefix :: forall f a. Index f a -> TermList f
minSize_ :: forall f a. Index f a -> Int
..} =
  Bool
nonEmpty Bool -> Bool -> Bool
&&
  Bool
noNilVars Bool -> Bool -> Bool
&&
  Bool
maxPrefix Bool -> Bool -> Bool
&&
  Bool
sizeCorrect Bool -> Bool -> Bool
&&
  forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall f a. Index f a -> Bool
invariant (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd (forall a. Array a -> [(Int, a)]
toList Array (Index f a)
fun)) Bool -> Bool -> Bool
&&
  forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall f a. Index f a -> Bool
invariant (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd (forall a. Numbered a -> [(Int, a)]
Numbered.toList Numbered (Index f a)
var))
  where
    nonEmpty :: Bool
nonEmpty = -- Index should not be empty
      Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
List.null [a]
here) Bool -> Bool -> Bool
||
      Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
List.null (forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall f a. Index f a -> Bool
null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) (forall a. Array a -> [(Int, a)]
toList Array (Index f a)
fun))) Bool -> Bool -> Bool
||
      Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
List.null (forall a. Numbered a -> [(Int, a)]
Numbered.toList Numbered (Index f a)
var))
    noNilVars :: Bool
noNilVars = -- the var field should not contain any Nils
      forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall f a. Index f a -> Bool
null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) (forall a. Numbered a -> [(Int, a)]
Numbered.toList Numbered (Index f a)
var)
    maxPrefix :: Bool
maxPrefix -- prefix should be used if possible
      | forall (t :: * -> *) a. Foldable t => t a -> Bool
List.null [a]
here =
        forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall f a. Index f a -> Bool
null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) (forall a. Array a -> [(Int, a)]
toList Array (Index f a)
fun)) forall a. Num a => a -> a -> a
+
        forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall a. Numbered a -> [(Int, a)]
Numbered.toList Numbered (Index f a)
var) forall a. Ord a => a -> a -> Bool
> Int
1
      | Bool
otherwise = Bool
True
    sizeCorrect :: Bool
sizeCorrect -- size field must be correct
      | forall (t :: * -> *) a. Foldable t => t a -> Bool
List.null [a]
here =
        (Int
minSize_ forall a. Num a => a -> a -> a
- forall f. TermList f -> Int
lenList TermList f
prefix forall a. Num a => a -> a -> a
- Int
1) forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem`
        forall a b. (a -> b) -> [a] -> [b]
map (forall f a. Index f a -> Int
minSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) (forall a. Array a -> [(Int, a)]
toList Array (Index f a)
fun) forall a. [a] -> [a] -> [a]
++
        forall a b. (a -> b) -> [a] -> [b]
map (forall f a. Index f a -> Int
minSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) (forall a. Numbered a -> [(Int, a)]
Numbered.toList Numbered (Index f a)
var)
      | Bool
otherwise =
        Int
minSize_ forall a. Eq a => a -> a -> Bool
== forall f. TermList f -> Int
lenList TermList f
prefix

instance Default (Index f a) where def :: Index f a
def = forall f a. Index f a
Nil

-- | An empty index.
empty :: Index f a
empty :: forall f a. Index f a
empty = forall f a. Index f a
Nil

-- | Is the index empty?
null :: Index f a -> Bool
null :: forall f a. Index f a -> Bool
null Index f a
Nil = Bool
True
null Index f a
_ = Bool
False

-- | An index with one entry.
singleton :: Term f -> a -> Index f a
singleton :: forall f a. Term f -> a -> Index f a
singleton !Term f
t a
x = forall f a. TermList f -> [a] -> Index f a
leaf (forall f. Term f -> TermList f
Term.singleton Term f
t) [a
x]

-- A leaf node, perhaps with a prefix.
leaf :: TermList f -> [a] -> Index f a
leaf :: forall f a. TermList f -> [a] -> Index f a
leaf !TermList f
_ [] = forall f a. Index f a
Nil
leaf TermList f
t [a]
xs = forall f a.
Int
-> TermList f
-> [a]
-> Array (Index f a)
-> Numbered (Index f a)
-> Index f a
Index (forall f. TermList f -> Int
lenList TermList f
t) TermList f
t [a]
xs forall a. Array a
newArray forall a. Numbered a
Numbered.empty

-- Add a prefix (given as a list of symbols) to all terms in an index.
addPrefix :: [Term f] -> Index f a -> Index f a
addPrefix :: forall f a. [Term f] -> Index f a -> Index f a
addPrefix [Term f]
_ Index f a
Nil = forall f a. Index f a
Nil
addPrefix [] Index f a
idx = Index f a
idx
addPrefix [Term f]
ts Index f a
idx =
  Index f a
idx {
    minSize_ :: Int
minSize_ = forall f a. Index f a -> Int
minSize_ Index f a
idx forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. Foldable t => t a -> Int
length [Term f]
ts,
    prefix :: TermList f
prefix = forall a. Build a => a -> TermList (BuildFun a)
buildList (forall a. Monoid a => [a] -> a
mconcat (forall a b. (a -> b) -> [a] -> [b]
map forall {f}. Term f -> Builder f
atom [Term f]
ts) forall a. Monoid a => a -> a -> a
`mappend` forall a. Build a => a -> Builder (BuildFun a)
builder (forall f a. Index f a -> TermList f
prefix Index f a
idx)) }
  where
    atom :: Term f -> Builder f
atom (Var Var
x) = forall f. Var -> Builder f
Term.var Var
x
    atom (App Fun f
f TermList f
_) = forall f. Fun f -> Builder f
con Fun f
f

-- Smart constructor for Index.
index :: [a] -> Array (Index f a) -> Numbered (Index f a) -> Index f a
index :: forall a f.
[a] -> Array (Index f a) -> Numbered (Index f a) -> Index f a
index [a]
here Array (Index f a)
fun Numbered (Index f a)
var =
  case ([a]
here, [(Int, Index f a)]
fun', forall a. Numbered a -> [(Int, a)]
Numbered.toList Numbered (Index f a)
var') of
    ([], [], []) ->
      forall f a. Index f a
Nil
    ([], [(Int
f, Index f a
idx)], []) ->
      Index f a
idx{minSize_ :: Int
minSize_ = forall a. Enum a => a -> a
succ (forall f a. Index f a -> Int
minSize_ Index f a
idx),
          prefix :: TermList f
prefix = forall a. Build a => a -> TermList (BuildFun a)
buildList (forall f. Fun f -> Builder f
con (forall f. Int -> Fun f
Core.F Int
f) forall a. Monoid a => a -> a -> a
`mappend` forall a. Build a => a -> Builder (BuildFun a)
builder (forall f a. Index f a -> TermList f
prefix Index f a
idx))}
    ([], [], [(Int
x, Index f a
idx)]) ->
      Index f a
idx{minSize_ :: Int
minSize_ = forall a. Enum a => a -> a
succ (forall f a. Index f a -> Int
minSize_ Index f a
idx),
          prefix :: TermList f
prefix = forall a. Build a => a -> TermList (BuildFun a)
buildList (forall f. Var -> Builder f
Term.var (Int -> Var
V Int
x) forall a. Monoid a => a -> a -> a
`mappend` forall a. Build a => a -> Builder (BuildFun a)
builder (forall f a. Index f a -> TermList f
prefix Index f a
idx))}
    ([a], [(Int, Index f a)], [(Int, Index f a)])
_ ->
      Index {
        minSize_ :: Int
minSize_ = Int
size,
        prefix :: TermList f
prefix = forall f. TermList f
Term.empty,
        here :: [a]
here = [a]
here,
        fun :: Array (Index f a)
fun = Array (Index f a)
fun,
        var :: Numbered (Index f a)
var = Numbered (Index f a)
var' }
  where
    var' :: Numbered (Index f a)
var' = forall a. (a -> Bool) -> Numbered a -> Numbered a
Numbered.filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall f a. Index f a -> Bool
null) Numbered (Index f a)
var
    fun' :: [(Int, Index f a)]
fun' = forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall f a. Index f a -> Bool
null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) (forall a. Array a -> [(Int, a)]
toList Array (Index f a)
fun)
    size :: Int
size =
      forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum forall a b. (a -> b) -> a -> b
$
        [Int
0 | Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
List.null [a]
here)] forall a. [a] -> [a] -> [a]
++
        forall a b. (a -> b) -> [a] -> [b]
map (forall a. Enum a => a -> a
succ forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall f a. Index f a -> Int
minSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Int, Index f a)]
fun' forall a. [a] -> [a] -> [a]
++
        forall a b. (a -> b) -> [a] -> [b]
map (forall a. Enum a => a -> a
succ forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall f a. Index f a -> Int
minSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) (forall a. Numbered a -> [(Int, a)]
Numbered.toList Numbered (Index f a)
var')

-- | Insert an entry into the index.
insert :: (Symbolic a, ConstantOf a ~ f) => Term f -> a -> Index f a -> Index f a
insert :: forall a f.
(Symbolic a, ConstantOf a ~ f) =>
Term f -> a -> Index f a -> Index f a
insert = forall a f.
(Symbolic a, ConstantOf a ~ f) =>
(a -> [a] -> [a]) -> Term f -> a -> Index f a -> Index f a
modify (:)

-- | Delete an entry from the index.
delete :: (Eq a, Symbolic a, ConstantOf a ~ f) => Term f -> a -> Index f a -> Index f a
delete :: forall a f.
(Eq a, Symbolic a, ConstantOf a ~ f) =>
Term f -> a -> Index f a -> Index f a
delete =
  forall a f.
(Symbolic a, ConstantOf a ~ f) =>
(a -> [a] -> [a]) -> Term f -> a -> Index f a -> Index f a
modify forall a b. (a -> b) -> a -> b
$ \a
x [a]
xs ->
    if a
x forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`List.elem` [a]
xs then forall a. Eq a => a -> [a] -> [a]
List.delete a
x [a]
xs
    else forall a. HasCallStack => String -> a
error String
"deleted term not found in index"

-- General-purpose function for modifying the index.
modify :: (Symbolic a, ConstantOf a ~ f) =>
  (a -> [a] -> [a]) ->
  Term f -> a -> Index f a -> Index f a
modify :: forall a f.
(Symbolic a, ConstantOf a ~ f) =>
(a -> [a] -> [a]) -> Term f -> a -> Index f a -> Index f a
modify a -> [a] -> [a]
f !Term f
t0 !a
v0 !Index f a
idx = [Term f] -> TermList f -> Index f a -> Index f a
aux [] (forall f. Term f -> TermList f
Term.singleton Term f
t) Index f a
idx
  where
    (!Term f
t, !a
v) = forall a. Symbolic a => a -> a
canonicalise (Term f
t0, a
v0) 

    aux :: [Term f] -> TermList f -> Index f a -> Index f a
aux [] TermList f
t Index f a
Nil =
      forall f a. TermList f -> [a] -> Index f a
leaf TermList f
t (a -> [a] -> [a]
f a
v [])

    -- Non-empty prefix
    aux [Term f]
syms (ConsSym{hd :: forall f. TermList f -> Term f
hd = t :: Term f
t@(Var Var
x), rest :: forall f. TermList f -> TermList f
rest = TermList f
ts})
      idx :: Index f a
idx@Index{prefix :: forall f a. Index f a -> TermList f
prefix = ConsSym{hd :: forall f. TermList f -> Term f
hd = Var Var
y, rest :: forall f. TermList f -> TermList f
rest = TermList f
us}}
      | Var
x forall a. Eq a => a -> a -> Bool
== Var
y =
        [Term f] -> TermList f -> Index f a -> Index f a
aux (Term f
tforall a. a -> [a] -> [a]
:[Term f]
syms) TermList f
ts Index f a
idx{prefix :: TermList f
prefix = TermList f
us, minSize_ :: Int
minSize_ = forall f a. Index f a -> Int
minSize_ Index f a
idxforall a. Num a => a -> a -> a
-Int
1}
    aux [Term f]
syms (ConsSym{hd :: forall f. TermList f -> Term f
hd = t :: Term f
t@(App Fun f
f TermList f
_), rest :: forall f. TermList f -> TermList f
rest = TermList f
ts})
      idx :: Index f a
idx@Index{prefix :: forall f a. Index f a -> TermList f
prefix = ConsSym{hd :: forall f. TermList f -> Term f
hd = App Fun f
g TermList f
_, rest :: forall f. TermList f -> TermList f
rest = TermList f
us}}
      | Fun f
f forall a. Eq a => a -> a -> Bool
== Fun f
g =
        [Term f] -> TermList f -> Index f a -> Index f a
aux (Term f
tforall a. a -> [a] -> [a]
:[Term f]
syms) TermList f
ts Index f a
idx{prefix :: TermList f
prefix = TermList f
us, minSize_ :: Int
minSize_ = forall f a. Index f a -> Int
minSize_ Index f a
idxforall a. Num a => a -> a -> a
-Int
1}
    aux [] TermList f
t idx :: Index f a
idx@Index{prefix :: forall f a. Index f a -> TermList f
prefix = Cons{}} =
      [Term f] -> TermList f -> Index f a -> Index f a
aux [] TermList f
t (forall f a. Index f a -> Index f a
expand Index f a
idx)
    aux syms :: [Term f]
syms@(Term f
_:[Term f]
_) TermList f
t Index f a
idx =
      forall f a. [Term f] -> Index f a -> Index f a
addPrefix (forall a. [a] -> [a]
reverse [Term f]
syms) forall a b. (a -> b) -> a -> b
$ [Term f] -> TermList f -> Index f a -> Index f a
aux [] TermList f
t Index f a
idx

    -- Empty prefix
    aux [] TermList f
Empty Index f a
idx =
      forall a f.
[a] -> Array (Index f a) -> Numbered (Index f a) -> Index f a
index (a -> [a] -> [a]
f a
v (forall f a. Index f a -> [a]
here Index f a
idx)) (forall f a. Index f a -> Array (Index f a)
fun Index f a
idx) (forall f a. Index f a -> Numbered (Index f a)
var Index f a
idx)
    aux [] ConsSym{hd :: forall f. TermList f -> Term f
hd = App Fun f
f TermList f
_, rest :: forall f. TermList f -> TermList f
rest = TermList f
u} Index f a
idx =
      forall a f.
[a] -> Array (Index f a) -> Numbered (Index f a) -> Index f a
index (forall f a. Index f a -> [a]
here Index f a
idx)
        (forall a. Default a => Int -> a -> Array a -> Array a
update (forall f. Fun f -> Int
fun_id Fun f
f) Index f a
idx' (forall f a. Index f a -> Array (Index f a)
fun Index f a
idx))
        (forall f a. Index f a -> Numbered (Index f a)
var Index f a
idx)
      where
        idx' :: Index f a
idx' = [Term f] -> TermList f -> Index f a -> Index f a
aux [] TermList f
u (forall f a. Index f a -> Array (Index f a)
fun Index f a
idx forall a. Default a => Array a -> Int -> a
! forall f. Fun f -> Int
fun_id Fun f
f)
    aux [] ConsSym{hd :: forall f. TermList f -> Term f
hd = Var Var
x, rest :: forall f. TermList f -> TermList f
rest = TermList f
u} Index f a
idx =
      forall a f.
[a] -> Array (Index f a) -> Numbered (Index f a) -> Index f a
index (forall f a. Index f a -> [a]
here Index f a
idx) (forall f a. Index f a -> Array (Index f a)
fun Index f a
idx)
        (forall a. Int -> a -> (a -> a) -> Numbered a -> Numbered a
Numbered.modify (Var -> Int
var_id Var
x) forall f a. Index f a
Nil ([Term f] -> TermList f -> Index f a -> Index f a
aux [] TermList f
u) (forall f a. Index f a -> Numbered (Index f a)
var Index f a
idx))

-- Helper for modify:
-- Take an index with a prefix and pull out the first symbol of the prefix,
-- giving an index which doesn't start with a prefix.
{-# INLINE expand #-}
expand :: Index f a -> Index f a
expand :: forall f a. Index f a -> Index f a
expand idx :: Index f a
idx@Index{minSize_ :: forall f a. Index f a -> Int
minSize_ = Int
size, prefix :: forall f a. Index f a -> TermList f
prefix = ConsSym{hd :: forall f. TermList f -> Term f
hd = Term f
t, rest :: forall f. TermList f -> TermList f
rest = TermList f
ts}} =
  case Term f
t of
    Var Var
x ->
      Index {
        minSize_ :: Int
minSize_ = Int
size,
        prefix :: TermList f
prefix = forall f. TermList f
Term.empty,
        here :: [a]
here = [],
        fun :: Array (Index f a)
fun = forall a. Array a
newArray,
        var :: Numbered (Index f a)
var = forall a. Int -> a -> Numbered a
Numbered.singleton (Var -> Int
var_id Var
x) Index f a
idx { prefix :: TermList f
prefix = TermList f
ts, minSize_ :: Int
minSize_ = Int
size forall a. Num a => a -> a -> a
- Int
1 }}
    App Fun f
f TermList f
_ ->
      Index {
        minSize_ :: Int
minSize_ = Int
size,
        prefix :: TermList f
prefix = forall f. TermList f
Term.empty,
        here :: [a]
here = [],
        fun :: Array (Index f a)
fun = forall a. Default a => Int -> a -> Array a
Array.singleton (forall f. Fun f -> Int
fun_id Fun f
f) Index f a
idx { prefix :: TermList f
prefix = TermList f
ts, minSize_ :: Int
minSize_ = Int
size forall a. Num a => a -> a -> a
- Int
1 },
        var :: Numbered (Index f a)
var = forall a. Numbered a
Numbered.empty }

-- | Look up a term in the index. Finds all key-value such that the search term
-- is an instance of the key, and returns an instance of the the value which
-- makes the search term exactly equal to the key.
{-# INLINE lookup #-}
lookup :: (Has a b, Symbolic b, Has b (TermOf b)) => TermOf b -> Index (ConstantOf b) a -> [b]
lookup :: forall a b.
(Has a b, Symbolic b, Has b (TermOf b)) =>
TermOf b -> Index (ConstantOf b) a -> [b]
lookup TermOf b
t Index (ConstantOf b) a
idx = forall a b.
(Has a b, Symbolic b, Has b (TermOf b)) =>
TermListOf b -> Index (ConstantOf b) a -> [b]
lookupList (forall f. Term f -> TermList f
Term.singleton TermOf b
t) Index (ConstantOf b) a
idx

{-# INLINEABLE lookupList #-}
lookupList :: (Has a b, Symbolic b, Has b (TermOf b)) => TermListOf b -> Index (ConstantOf b) a -> [b]
lookupList :: forall a b.
(Has a b, Symbolic b, Has b (TermOf b)) =>
TermListOf b -> Index (ConstantOf b) a -> [b]
lookupList TermListOf b
t Index (ConstantOf b) a
idx =
  [ forall a s.
(Symbolic a, Substitution s, SubstFun s ~ ConstantOf a) =>
s -> a -> a
subst Subst (ConstantOf b)
sub (forall a b. Has a b => a -> b
the a
x)
  | (Subst (ConstantOf b)
sub, a
x) <- forall f a. TermList f -> Index f a -> [(Subst f, a)]
matchesList TermListOf b
t Index (ConstantOf b) a
idx ]

-- | Look up a term in the index. Like 'lookup', but returns the exact value
-- that was inserted into the index, not an instance. Also returns a substitution
-- which when applied to the value gives you the matching instance.
{-# INLINE matches #-}
matches :: Term f -> Index f a -> [(Subst f, a)]
matches :: forall f a. Term f -> Index f a -> [(Subst f, a)]
matches Term f
t Index f a
idx = forall f a. TermList f -> Index f a -> [(Subst f, a)]
matchesList (forall f. Term f -> TermList f
Term.singleton Term f
t) Index f a
idx

matchesList :: TermList f -> Index f a -> [(Subst f, a)]
matchesList :: forall f a. TermList f -> Index f a -> [(Subst f, a)]
matchesList TermList f
t Index f a
idx =
  forall f a. Stack f a -> [(Subst f, a)]
run (forall f a.
TermList f -> Bindings f -> Index f a -> Stack f a -> Stack f a
search TermList f
t forall f. Bindings f
emptyBindings Index f a
idx forall f a. Stack f a
Stop)

-- | Return all elements of the index.
elems :: Index f a -> [a]
elems :: forall f a. Index f a -> [a]
elems Index f a
Nil = []
elems Index f a
idx =
  forall f a. Index f a -> [a]
here Index f a
idx forall a. [a] -> [a] -> [a]
++
  forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall f a. Index f a -> [a]
elems (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd (forall a. Array a -> [(Int, a)]
toList (forall f a. Index f a -> Array (Index f a)
fun Index f a
idx))) forall a. [a] -> [a] -> [a]
++
  forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall f a. Index f a -> [a]
elems (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd (forall a. Numbered a -> [(Int, a)]
Numbered.toList (forall f a. Index f a -> Numbered (Index f a)
var Index f a
idx)))

-- | Create an index from a list of items
fromList :: (Symbolic a, ConstantOf a ~ f) => [(Term f, a)] -> Index f a
fromList :: forall a f.
(Symbolic a, ConstantOf a ~ f) =>
[(Term f, a)] -> Index f a
fromList [(Term f, a)]
xs = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a f.
(Symbolic a, ConstantOf a ~ f) =>
Term f -> a -> Index f a -> Index f a
insert) forall f a. Index f a
empty [(Term f, a)]
xs

-- | Create an index from a list of items
fromListWith :: (Symbolic a, ConstantOf a ~ f) => (a -> Term f) -> [a] -> Index f a
fromListWith :: forall a f.
(Symbolic a, ConstantOf a ~ f) =>
(a -> Term f) -> [a] -> Index f a
fromListWith a -> Term f
f [a]
xs = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
x -> forall a f.
(Symbolic a, ConstantOf a ~ f) =>
Term f -> a -> Index f a -> Index f a
insert (a -> Term f
f a
x) a
x) forall f a. Index f a
empty [a]
xs

----------------------------------------------------------------------
-- Substitutions used internally during lookup.
----------------------------------------------------------------------

-- We represent a substitution as a list of terms, in
-- reverse order. That is, the substitution
-- {x1->t1, ..., xn->tn} is represented as
-- [xn, x{n-1}, ..., x1].
data Bindings f =
  Bindings
    {-# UNPACK #-} !Int -- the highest-numbered variable (n)
    !(BindList f)       -- the list of terms ([xn, ..., x1])

data BindList f = NilB | ConsB {-# UNPACK #-} !(TermList f) !(BindList f)

{-# INLINE emptyBindings #-}
-- An empty substitution
emptyBindings :: Bindings f
emptyBindings :: forall f. Bindings f
emptyBindings = forall f. Int -> BindList f -> Bindings f
Bindings (-Int
1) forall f. BindList f
NilB

{-# INLINE extendBindings #-}
-- Extend a substitution.
-- The variable bound must either be present in the substitution,
-- or the current highest-numbered variable plus one.
extendBindings :: Var -> Term f -> Bindings f -> Maybe (Bindings f)
extendBindings :: forall f. Var -> Term f -> Bindings f -> Maybe (Bindings f)
extendBindings (V Int
x) Term f
t (Bindings Int
n BindList f
bs)
  | Int
x forall a. Ord a => a -> a -> Bool
> Int
n = forall a. a -> Maybe a
Just (forall f. Int -> BindList f -> Bindings f
Bindings (Int
nforall a. Num a => a -> a -> a
+Int
1) (forall f. TermList f -> BindList f -> BindList f
ConsB (forall f. Term f -> TermList f
Term.singleton Term f
t) BindList f
bs))
  | BindList f
bs forall f. BindList f -> Int -> TermList f
`at` (Int
n forall a. Num a => a -> a -> a
- Int
x) forall a. Eq a => a -> a -> Bool
== forall f. Term f -> TermList f
Term.singleton Term f
t = forall a. a -> Maybe a
Just (forall f. Int -> BindList f -> Bindings f
Bindings Int
n BindList f
bs)
  | Bool
otherwise = forall a. Maybe a
Nothing

at :: BindList f -> Int -> TermList f
at :: forall f. BindList f -> Int -> TermList f
at (ConsB TermList f
t BindList f
_) Int
0 = TermList f
t
at (ConsB TermList f
_ BindList f
b) Int
n = forall f. BindList f -> Int -> TermList f
at BindList f
b (Int
nforall a. Num a => a -> a -> a
-Int
1)

-- Convert a substitution into an ordinary Subst.
toSubst :: Bindings f -> Subst f
toSubst :: forall f. Bindings f -> Subst f
toSubst (Bindings Int
n BindList f
bs) =
  forall f. IntMap (TermList f) -> Subst f
Subst (forall a. [(Int, a)] -> IntMap a
IntMap.fromDistinctAscList (forall {a} {f}.
Num a =>
a -> BindList f -> [(a, TermList f)] -> [(a, TermList f)]
loop Int
n BindList f
bs []))
  where
    loop :: a -> BindList f -> [(a, TermList f)] -> [(a, TermList f)]
loop !a
_ !BindList f
_ ![(a, TermList f)]
_ | Bool
False = forall a. HasCallStack => a
undefined
    loop a
_ BindList f
NilB [(a, TermList f)]
sub = [(a, TermList f)]
sub
    loop a
n (ConsB TermList f
t BindList f
bs) [(a, TermList f)]
sub =
      a -> BindList f -> [(a, TermList f)] -> [(a, TermList f)]
loop (a
nforall a. Num a => a -> a -> a
-a
1) BindList f
bs ((a
n, TermList f
t)forall a. a -> [a] -> [a]
:[(a, TermList f)]
sub)

----------------------------------------------------------------------
-- Lookup.
----------------------------------------------------------------------

-- To get predictable performance, lookup uses an explicit stack
-- instead of a lazy list to control backtracking.
data Stack f a =
  -- We only ever backtrack into variable edges, not function edges.
  -- This stack frame represents a search of the variable edges of a
  -- node, starting at a particular variable.
  Frame {
    -- The term which should match the variable
    forall f a. Stack f a -> Term f
frame_term    :: {-# UNPACK #-} !(Term f),
    -- The remainder of the search term
    forall f a. Stack f a -> TermList f
frame_terms   :: {-# UNPACK #-} !(TermList f),
    -- The current substitution
    forall f a. Stack f a -> Bindings f
frame_bind    :: {-# UNPACK #-} !(Bindings f),
    -- The list of variable edges
    forall f a. Stack f a -> Numbered (Index f a)
frame_indexes :: {-# UNPACK #-} !(Numbered (Index f a)),
    -- The starting variable number
    forall f a. Stack f a -> Int
frame_var     :: {-# UNPACK #-} !Int,
    -- The rest of the stack
    forall f a. Stack f a -> Stack f a
frame_rest    :: !(Stack f a) } |
  -- A stack frame which is used when we have found a matching node.
  Yield {
    -- The list of values found at this node
    forall f a. Stack f a -> [a]
yield_found  :: [a],
    -- The current substitution
    forall f a. Stack f a -> Bindings f
yield_binds  :: {-# UNPACK #-} !(Bindings f),
    -- The rest of the stack
    forall f a. Stack f a -> Stack f a
yield_rest   :: !(Stack f a) }
  -- End of stack.
  | Stop

-- Turn a stack into a list of results.
run :: Stack f a -> [(Subst f, a)]
run :: forall f a. Stack f a -> [(Subst f, a)]
run Stack f a
stack = forall symbol a. symbol -> a -> a
stamp String
"index lookup" (forall f a. Stack f a -> [(Subst f, a)]
run1 Stack f a
stack) 
  where
    run1 :: Stack f a -> [(Subst f, a)]
run1 Stack f a
Stop = []
    run1 Frame{Int
Numbered (Index f a)
Term f
TermList f
Stack f a
Bindings f
frame_rest :: Stack f a
frame_var :: Int
frame_indexes :: Numbered (Index f a)
frame_bind :: Bindings f
frame_terms :: TermList f
frame_term :: Term f
frame_rest :: forall f a. Stack f a -> Stack f a
frame_var :: forall f a. Stack f a -> Int
frame_indexes :: forall f a. Stack f a -> Numbered (Index f a)
frame_bind :: forall f a. Stack f a -> Bindings f
frame_terms :: forall f a. Stack f a -> TermList f
frame_term :: forall f a. Stack f a -> Term f
..} =
      Stack f a -> [(Subst f, a)]
run1 (forall f a.
Term f
-> TermList f
-> Bindings f
-> Numbered (Index f a)
-> Int
-> Stack f a
-> Stack f a
searchVars Term f
frame_term TermList f
frame_terms Bindings f
frame_bind Numbered (Index f a)
frame_indexes Int
frame_var Stack f a
frame_rest)
    run1 Yield{[a]
Stack f a
Bindings f
yield_rest :: Stack f a
yield_binds :: Bindings f
yield_found :: [a]
yield_rest :: forall f a. Stack f a -> Stack f a
yield_binds :: forall f a. Stack f a -> Bindings f
yield_found :: forall f a. Stack f a -> [a]
..} =
      forall a b. (a -> b) -> [a] -> [b]
map (forall f. Bindings f -> Subst f
toSubst Bindings f
yield_binds,) [a]
yield_found forall a. [a] -> [a] -> [a]
++ forall f a. Stack f a -> [(Subst f, a)]
run Stack f a
yield_rest

-- Search starting with a given substitution.
{-# INLINE search #-}
search :: TermList f -> Bindings f -> Index f a -> Stack f a -> Stack f a
search :: forall f a.
TermList f -> Bindings f -> Index f a -> Stack f a -> Stack f a
search !TermList f
_ !Bindings f
_ !Index f a
_ Stack f a
_ | Bool
False = forall a. HasCallStack => a
undefined
search TermList f
t Bindings f
binds Index f a
idx Stack f a
rest =
  case Index f a
idx of
    Index f a
Nil -> Stack f a
rest
    Index{Int
[a]
Array (Index f a)
Numbered (Index f a)
TermList f
var :: Numbered (Index f a)
fun :: Array (Index f a)
here :: [a]
prefix :: TermList f
minSize_ :: Int
var :: forall f a. Index f a -> Numbered (Index f a)
fun :: forall f a. Index f a -> Array (Index f a)
here :: forall f a. Index f a -> [a]
prefix :: forall f a. Index f a -> TermList f
minSize_ :: forall f a. Index f a -> Int
..}
      | forall f. TermList f -> Int
lenList TermList f
t forall a. Ord a => a -> a -> Bool
< forall f a. Index f a -> Int
minSize Index f a
idx ->
        Stack f a
rest -- the search term is smaller than any in this index
      | Bool
otherwise ->
        forall f a.
Bindings f
-> TermList f
-> TermList f
-> [a]
-> Array (Index f a)
-> Numbered (Index f a)
-> Stack f a
-> Stack f a
searchLoop Bindings f
binds TermList f
t TermList f
prefix [a]
here Array (Index f a)
fun Numbered (Index f a)
var Stack f a
rest

-- The main work of 'search' goes on here.
-- It is carefully tweaked to generate nice code,
-- in particular casing on each term list exactly once.
searchLoop ::
  -- Search term and substitution
  Bindings f -> TermList f ->
  -- Contents of the relevant node of the index
  TermList f -> [a] -> Array (Index f a) -> Numbered (Index f a) ->
  Stack f a -> Stack f a
searchLoop :: forall f a.
Bindings f
-> TermList f
-> TermList f
-> [a]
-> Array (Index f a)
-> Numbered (Index f a)
-> Stack f a
-> Stack f a
searchLoop !Bindings f
_ !TermList f
_ !TermList f
_ [a]
_ !Array (Index f a)
_ !Numbered (Index f a)
_ Stack f a
_ | Bool
False = forall a. HasCallStack => a
undefined
searchLoop Bindings f
binds TermList f
t TermList f
prefix [a]
here Array (Index f a)
fun Numbered (Index f a)
var Stack f a
rest =
  case TermList f
t of
    ConsSym{hd :: forall f. TermList f -> Term f
hd = Term f
thd, tl :: forall f. TermList f -> TermList f
tl = TermList f
ttl, rest :: forall f. TermList f -> TermList f
rest = TermList f
trest} ->
      case TermList f
prefix of
        ConsSym{hd :: forall f. TermList f -> Term f
hd = Term f
phd, tl :: forall f. TermList f -> TermList f
tl = TermList f
ptl, rest :: forall f. TermList f -> TermList f
rest = TermList f
prest} ->
          -- Check the search term against the prefix.
          case (Term f
thd, Term f
phd) of
            (Term f
_, Var Var
x) ->
              case forall f. Var -> Term f -> Bindings f -> Maybe (Bindings f)
extendBindings Var
x Term f
thd Bindings f
binds of
                Just Bindings f
binds' ->
                  forall f a.
Bindings f
-> TermList f
-> TermList f
-> [a]
-> Array (Index f a)
-> Numbered (Index f a)
-> Stack f a
-> Stack f a
searchLoop Bindings f
binds' TermList f
ttl TermList f
ptl [a]
here Array (Index f a)
fun Numbered (Index f a)
var Stack f a
rest
                Maybe (Bindings f)
Nothing ->
                  Stack f a
rest
            (App Fun f
f TermList f
_, App Fun f
g TermList f
_) | Fun f
f forall a. Eq a => a -> a -> Bool
== Fun f
g ->
               -- Term and prefix start with same symbol, chop them off.
               forall f a.
Bindings f
-> TermList f
-> TermList f
-> [a]
-> Array (Index f a)
-> Numbered (Index f a)
-> Stack f a
-> Stack f a
searchLoop Bindings f
binds TermList f
trest TermList f
prest [a]
here Array (Index f a)
fun Numbered (Index f a)
var Stack f a
rest
            (Term f, Term f)
_ ->
              -- Term and prefix don't match.
              Stack f a
rest
        TermList f
_ ->
          -- We've exhausted the prefix, so let's descend into the tree.
          -- Seems to work better to explore the function node first.
          case Term f
thd of
            App Fun f
f TermList f
_ | idx :: Index f a
idx@Index{} <- Array (Index f a)
fun forall a. Default a => Array a -> Int -> a
! forall f. Fun f -> Int
fun_id Fun f
f ->
              -- Avoid creating a frame unnecessarily.
              case forall a. Numbered a -> Int
Numbered.size Numbered (Index f a)
var of
                Int
0 -> forall f a.
TermList f -> Bindings f -> Index f a -> Stack f a -> Stack f a
search TermList f
trest Bindings f
binds Index f a
idx Stack f a
rest
                Int
_ -> forall f a.
TermList f -> Bindings f -> Index f a -> Stack f a -> Stack f a
search TermList f
trest Bindings f
binds Index f a
idx forall a b. (a -> b) -> a -> b
$! forall f a.
Term f
-> TermList f
-> Bindings f
-> Numbered (Index f a)
-> Int
-> Stack f a
-> Stack f a
Frame Term f
thd TermList f
ttl Bindings f
binds Numbered (Index f a)
var Int
0 Stack f a
rest
            Term f
_ -> -- no function match
              case forall a. Numbered a -> Int
Numbered.size Numbered (Index f a)
var of
                Int
0 -> Stack f a
rest
                Int
_ -> forall f a.
Term f
-> TermList f
-> Bindings f
-> Numbered (Index f a)
-> Int
-> Stack f a
-> Stack f a
searchVars Term f
thd TermList f
ttl Bindings f
binds Numbered (Index f a)
var Int
0 Stack f a
rest
    TermList f
_ ->
      case TermList f
prefix of
        TermList f
Empty ->
          -- The search term matches this node.
          case [a]
here of
            [] -> Stack f a
rest
            [a]
_ -> forall f a. [a] -> Bindings f -> Stack f a -> Stack f a
Yield [a]
here Bindings f
binds Stack f a
rest
        TermList f
_ ->
          -- We've run out of search term - it doesn't match this node.
          Stack f a
rest

-- Search the variable-labelled edges of a node,
-- starting with a particular variable.
searchVars ::
  -- Term (with head separate) and substitution
  Term f -> TermList f -> Bindings f ->
  -- Variable edges and starting variable
  Numbered (Index f a) -> Int ->
  Stack f a -> Stack f a
searchVars :: forall f a.
Term f
-> TermList f
-> Bindings f
-> Numbered (Index f a)
-> Int
-> Stack f a
-> Stack f a
searchVars !Term f
_ !TermList f
_ !Bindings f
_ !Numbered (Index f a)
_ !Int
_ Stack f a
_ | Bool
False = forall a. HasCallStack => a
undefined
searchVars Term f
t TermList f
ts Bindings f
binds Numbered (Index f a)
var Int
start Stack f a
rest
  | Int
start forall a. Ord a => a -> a -> Bool
>= forall a. Numbered a -> Int
Numbered.size Numbered (Index f a)
var = Stack f a
rest
  | Bool
otherwise =
    let (Int
x, Index f a
idx) = Numbered (Index f a)
var forall a. Numbered a -> Int -> (Int, a)
Numbered.! Int
start in
    case forall f. Var -> Term f -> Bindings f -> Maybe (Bindings f)
extendBindings (Int -> Var
V Int
x) Term f
t Bindings f
binds of
      Just Bindings f
binds' ->
        forall f a.
TermList f -> Bindings f -> Index f a -> Stack f a -> Stack f a
search TermList f
ts Bindings f
binds' Index f a
idx forall a b. (a -> b) -> a -> b
$!
        if Int
start forall a. Num a => a -> a -> a
+ Int
1 forall a. Eq a => a -> a -> Bool
== forall a. Numbered a -> Int
Numbered.size Numbered (Index f a)
var then Stack f a
rest
        else forall f a.
Term f
-> TermList f
-> Bindings f
-> Numbered (Index f a)
-> Int
-> Stack f a
-> Stack f a
Frame Term f
t TermList f
ts Bindings f
binds Numbered (Index f a)
var (Int
startforall a. Num a => a -> a -> a
+Int
1) Stack f a
rest
      Maybe (Bindings f)
Nothing ->
        forall f a.
Term f
-> TermList f
-> Bindings f
-> Numbered (Index f a)
-> Int
-> Stack f a
-> Stack f a
searchVars Term f
t TermList f
ts Bindings f
binds Numbered (Index f a)
var (Int
startforall a. Num a => a -> a -> a
+Int
1) Stack f a
rest