Prioritized critical bit tree data structure

Pcbt.hs

{-# Language TypeFamilies #-}

module Pcbt where

import Prelude hiding (lookup)
import Data.Set (Set)
import Data.Map (Map)
import qualified Data.Set as Set
import qualified Data.Map.Strict as Map
import Data.Maybe

data Bit = Zero | One | Star deriving (Show,Eq)
data Pcbt p = Miss | Hit | Branch { path :: p, maxPath' :: p, zero :: Pcbt p, one :: Pcbt p }
type Permutation p = p -> p

class Bits a where
  type Path a :: *
  bits :: a -> Path a -> Bit
  pcbt :: a -> Pcbt (Path a)

branch :: Ord p => p -> Pcbt p -> Pcbt p -> Pcbt p
branch _ Miss Miss = Miss
branch i a b = Branch i (maxPath i a `max` maxPath i b) a b

maxPath :: p -> Pcbt p -> p
maxPath _ (Branch _ p _ _) = p
maxPath p _ = p

isMiss :: Pcbt p -> Bool
isMiss Miss = True
isMiss _ = False

-- | Trim this `Pcbt p` to only contain hits against the given query.
trim :: Ord p => (p -> Bit) -> Pcbt p -> Pcbt p
trim _ Miss = Miss
trim _ Hit = Hit
trim query (Branch i _ zero one) = case query i of
  Zero -> branch i (trim query zero) Miss
  One -> branch i Miss (trim query one)
  Star -> branch i (trim query zero) (trim query one)

intersectOn :: Ord p => (p -> Maybe p) -> Pcbt p -> Pcbt p -> Pcbt p
intersectOn _ _ Miss = Miss
intersectOn _ a Hit = a
intersectOn _ Miss _ = Miss
intersectOn _ Hit _ = Hit
intersectOn on (Branch p maxP zero one) b2@(Branch pi maxPi zeroi onei) = case on p of
  Just p' | p' == pi -> branch p (intersectOn on zero zeroi) (intersectOn on one onei)
  _ -> branch p (intersectOn on zero b2) (intersectOn on one b2)

smallest :: Ord p => Pcbt p -> Pcbt p
smallest = smallestBy id

smallestBy :: Ord p => Permutation p -> Pcbt p -> Pcbt p
smallestBy permute t = fst $ go t Nothing where
  go Miss low = (Miss, low)
  go Hit low = (Hit, low)
  go (Branch i maxP z o) low = case low of
    -- No need to check this branch if we've already seen an earlier bit that's zero
    Just j | j < permute maxP -> (Miss, low)
    _ ->
      let
        (z', low') = go z low
        zlow = if not (isMiss z') then lower (permute i) low' else low'
        (o', olow) = go o zlow
      in
        if zlow /= olow -- o added something even smaller, so pick `o` branch
        then (branch i Miss o', olow)
        else (branch i z' Miss, olow) -- index i breaks the tie, choose `z` branch
  lower i a = Just $ fromMaybe i a `min` i

union :: Ord p => Pcbt p -> Pcbt p -> Pcbt p
union Miss p = p
union p Miss = p
union Hit _ = Hit
union _ Hit = Hit
union b1@(Branch i p1 o1 z1) b2@(Branch j p2 o2 z2)
  | i == j = Branch i (p1 `max` p2) (o1 `union` o2) (z1 `union` z2)
  | otherwise = Branch i (p1 `max` p2) (o1 `union` b2) (z1 `union` b2)

intersect :: Ord p => Pcbt p -> Pcbt p -> Pcbt p
intersect Miss _ = Miss
intersect _ Miss = Miss
intersect Hit p = p
intersect p Hit = p
intersect b1@(Branch i p1 o1 z1) b2@(Branch j p2 o2 z2)
  | i == j = branch i (o1 `intersect` o2) (z1 `intersect` z2)
  | otherwise = branch i (o1 `intersect` b2) (z1 `intersect` b2)

-- | Replace misses against the query with `mempty`, hits with `h`,
-- and combine results using `mappend`.
reduce :: Monoid a => a -> Pcbt p -> (p -> Bit) -> a
reduce _ Miss _ = mempty
reduce h Hit _ = h
reduce h (Branch i _ one zero) query = case query i of
  Zero -> reduce h zero query
  One -> reduce h one query
  Star -> reduce h zero query `mappend` reduce h one query

exists :: Pcbt p -> (p -> Bit) -> Bool
exists t q = unOr $ reduce (Or True) t q

forall :: Pcbt p -> (p -> Bit) -> Bool
forall t q = unAnd $ reduce (And True) t q

count :: Pcbt p -> (p -> Bit) -> Word
count t q = unSum $ reduce (Sum 1) t q

summarize :: Bit -> Bit -> Bit
summarize One One = One
summarize Zero Zero = Zero
summarize _ _ = Star

newtype Or = Or { unOr :: Bool }
instance Monoid Or where
  mempty = Or False
  mappend (Or a) (Or b) = Or (a || b)

newtype And = And { unAnd :: Bool }
instance Monoid And where
  mempty = And True
  mappend (And a) (And b) = And (a && b)

newtype Sum a = Sum { unSum :: a }
instance Num a => Monoid (Sum a) where
  mempty = Sum 0
  mappend (Sum a) (Sum b) = Sum (a + b)