full-text-search-0.2.2.3: src/Data/SearchEngine/TermBag.hs
{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving, MultiParamTypeClasses,
TypeFamilies #-}
{-# LANGUAGE CPP #-}
#if __GLASGOW_HASKELL__ >= 908
{-# OPTIONS_GHC -Wno-x-partial #-}
#endif
module Data.SearchEngine.TermBag (
TermId(TermId), TermCount,
TermBag,
size,
fromList,
toList,
elems,
termCount,
denseTable,
invariant
) where
import qualified Data.Vector.Unboxed as Vec
import qualified Data.Vector.Unboxed.Mutable as MVec
import qualified Data.Vector.Generic as GVec
import qualified Data.Vector.Generic.Mutable as GMVec
import Control.Monad.ST
import Control.Monad (liftM)
import qualified Data.Map as Map
import Data.Word (Word32, Word8)
import Data.Bits
import qualified Data.List as List
import Data.Function (on)
newtype TermId = TermId { unTermId :: Word32 }
deriving (Eq, Ord, Show, Enum)
instance Bounded TermId where
minBound = TermId 0
maxBound = TermId 0x00FFFFFF
data TermBag = TermBag !Int !(Vec.Vector TermIdAndCount)
deriving Show
-- We sneakily stuff both the TermId and the bag count into one 32bit word
type TermIdAndCount = Word32
type TermCount = Word8
-- Bottom 24 bits is the TermId, top 8 bits is the bag count
termIdAndCount :: TermId -> Int -> TermIdAndCount
termIdAndCount (TermId termid) freq =
(min (fromIntegral freq) 255 `shiftL` 24)
.|. (termid .&. 0x00FFFFFF)
getTermId :: TermIdAndCount -> TermId
getTermId word = TermId (word .&. 0x00FFFFFF)
getTermCount :: TermIdAndCount -> TermCount
getTermCount word = fromIntegral (word `shiftR` 24)
invariant :: TermBag -> Bool
invariant (TermBag _ vec) =
strictlyAscending (Vec.toList vec)
where
strictlyAscending (a:xs@(b:_)) = getTermId a < getTermId b
&& strictlyAscending xs
strictlyAscending _ = True
size :: TermBag -> Int
size (TermBag sz _) = sz
elems :: TermBag -> [TermId]
elems (TermBag _ vec) = map getTermId (Vec.toList vec)
toList :: TermBag -> [(TermId, TermCount)]
toList (TermBag _ vec) = [ (getTermId x, getTermCount x)
| x <- Vec.toList vec ]
termCount :: TermBag -> TermId -> TermCount
termCount (TermBag _ vec) =
binarySearch 0 (Vec.length vec - 1)
where
binarySearch :: Int -> Int -> TermId -> TermCount
binarySearch !a !b !key
| a > b = 0
| otherwise =
let mid = (a + b) `div` 2
tidAndCount = vec Vec.! mid
in case compare key (getTermId tidAndCount) of
LT -> binarySearch a (mid-1) key
EQ -> getTermCount tidAndCount
GT -> binarySearch (mid+1) b key
fromList :: [TermId] -> TermBag
fromList termids =
let bag = Map.fromListWith (+) [ (t, 1) | t <- termids ]
sz = Map.foldl' (+) 0 bag
vec = Vec.fromListN (Map.size bag)
[ termIdAndCount termid freq
| (termid, freq) <- Map.toAscList bag ]
in TermBag sz vec
-- | Given a bunch of term bags, merge them into a table for easier subsequent
-- processing. This is bascially a sparse to dense conversion. Missing entries
-- are filled in with 0. We represent the table as one vector for the
-- term ids and a 2d array for the counts.
--
-- Unfortunately vector does not directly support 2d arrays and array does
-- not make it easy to trim arrays.
--
denseTable :: [TermBag] -> (Vec.Vector TermId, Vec.Vector TermCount)
denseTable termbags =
(tids, tcts)
where
-- First merge the TermIds into one array
-- then make a linear pass to create the counts array
-- filling in 0s or the counts as we find them
!numBags = length termbags
!tids = unionsTermId termbags
!numTerms = Vec.length tids
!numCounts = numTerms * numBags
!tcts = Vec.create (do
out <- MVec.new numCounts
sequence_
[ writeMergedTermCounts tids bag out i
| (n, TermBag _ bag) <- zip [0..] termbags
, let i = n * numTerms ]
return out
)
writeMergedTermCounts :: Vec.Vector TermId -> Vec.Vector TermIdAndCount ->
MVec.MVector s TermCount -> Int -> ST s ()
writeMergedTermCounts xs0 ys0 !out i0 =
-- assume xs & ys are sorted, and ys contains a subset of xs
go xs0 ys0 i0
where
go !xs !ys !i
| Vec.null ys = MVec.set (MVec.slice i (Vec.length xs) out) 0
| Vec.null xs = return ()
| otherwise = let x = Vec.head xs
ytc = Vec.head ys
y = getTermId ytc
c = getTermCount ytc
in case x == y of
True -> do MVec.write out i c
go (Vec.tail xs) (Vec.tail ys) (i+1)
False -> do MVec.write out i 0
go (Vec.tail xs) ys (i+1)
-- | Given a set of term bags, form the set of TermIds
--
unionsTermId :: [TermBag] -> Vec.Vector TermId
unionsTermId tbs =
case List.sortBy (compare `on` bagVecLength) tbs of
[] -> Vec.empty
[TermBag _ xs] -> (Vec.map getTermId xs)
(x0:x1:xs) -> List.foldl' union3 (union2 x0 x1) xs
where
bagVecLength (TermBag _ vec) = Vec.length vec
union2 :: TermBag -> TermBag -> Vec.Vector TermId
union2 (TermBag _ xs) (TermBag _ ys) =
Vec.create (MVec.new sizeBound >>= writeMergedUnion2 xs ys)
where
sizeBound = Vec.length xs + Vec.length ys
writeMergedUnion2 :: Vec.Vector TermIdAndCount -> Vec.Vector TermIdAndCount ->
MVec.MVector s TermId -> ST s (MVec.MVector s TermId)
writeMergedUnion2 xs0 ys0 !out = do
i <- go xs0 ys0 0
return $! MVec.take i out
where
go !xs !ys !i
| Vec.null xs = do Vec.copy (MVec.slice i (Vec.length ys) out)
(Vec.map getTermId ys)
return (i + Vec.length ys)
| Vec.null ys = do Vec.copy (MVec.slice i (Vec.length xs) out)
(Vec.map getTermId xs)
return (i + Vec.length xs)
| otherwise = let x = getTermId (Vec.head xs)
y = getTermId (Vec.head ys)
in case compare x y of
GT -> do MVec.write out i y
go xs (Vec.tail ys) (i+1)
EQ -> do MVec.write out i x
go (Vec.tail xs) (Vec.tail ys) (i+1)
LT -> do MVec.write out i x
go (Vec.tail xs) ys (i+1)
union3 :: Vec.Vector TermId -> TermBag -> Vec.Vector TermId
union3 xs (TermBag _ ys) =
Vec.create (MVec.new sizeBound >>= writeMergedUnion3 xs ys)
where
sizeBound = Vec.length xs + Vec.length ys
writeMergedUnion3 :: Vec.Vector TermId -> Vec.Vector TermIdAndCount ->
MVec.MVector s TermId -> ST s (MVec.MVector s TermId)
writeMergedUnion3 xs0 ys0 !out = do
i <- go xs0 ys0 0
return $! MVec.take i out
where
go !xs !ys !i
| Vec.null xs = do Vec.copy (MVec.slice i (Vec.length ys) out)
(Vec.map getTermId ys)
return (i + Vec.length ys)
| Vec.null ys = do Vec.copy (MVec.slice i (Vec.length xs) out) xs
return (i + Vec.length xs)
| otherwise = let x = Vec.head xs
y = getTermId (Vec.head ys)
in case compare x y of
GT -> do MVec.write out i y
go xs (Vec.tail ys) (i+1)
EQ -> do MVec.write out i x
go (Vec.tail xs) (Vec.tail ys) (i+1)
LT -> do MVec.write out i x
go (Vec.tail xs) ys (i+1)
------------------------------------------------------------------------------
-- verbose Unbox instances
--
instance MVec.Unbox TermId
newtype instance MVec.MVector s TermId = MV_TermId (MVec.MVector s Word32)
instance GMVec.MVector MVec.MVector TermId where
basicLength (MV_TermId v) = GMVec.basicLength v
basicUnsafeSlice i l (MV_TermId v) = MV_TermId (GMVec.basicUnsafeSlice i l v)
basicUnsafeNew l = MV_TermId `liftM` GMVec.basicUnsafeNew l
basicInitialize (MV_TermId v) = GMVec.basicInitialize v
basicUnsafeReplicate l x = MV_TermId `liftM` GMVec.basicUnsafeReplicate l (unTermId x)
basicUnsafeRead (MV_TermId v) i = TermId `liftM` GMVec.basicUnsafeRead v i
basicUnsafeWrite (MV_TermId v) i x = GMVec.basicUnsafeWrite v i (unTermId x)
basicClear (MV_TermId v) = GMVec.basicClear v
basicSet (MV_TermId v) x = GMVec.basicSet v (unTermId x)
basicUnsafeGrow (MV_TermId v) l = MV_TermId `liftM` GMVec.basicUnsafeGrow v l
basicUnsafeCopy (MV_TermId v) (MV_TermId v') = GMVec.basicUnsafeCopy v v'
basicUnsafeMove (MV_TermId v) (MV_TermId v') = GMVec.basicUnsafeMove v v'
basicOverlaps (MV_TermId v) (MV_TermId v') = GMVec.basicOverlaps v v'
{-# INLINE basicLength #-}
{-# INLINE basicUnsafeSlice #-}
{-# INLINE basicOverlaps #-}
{-# INLINE basicUnsafeNew #-}
{-# INLINE basicInitialize #-}
{-# INLINE basicUnsafeReplicate #-}
{-# INLINE basicUnsafeRead #-}
{-# INLINE basicUnsafeWrite #-}
{-# INLINE basicClear #-}
{-# INLINE basicSet #-}
{-# INLINE basicUnsafeCopy #-}
{-# INLINE basicUnsafeMove #-}
{-# INLINE basicUnsafeGrow #-}
newtype instance Vec.Vector TermId = V_TermId (Vec.Vector Word32)
instance GVec.Vector Vec.Vector TermId where
basicUnsafeFreeze (MV_TermId mv) = V_TermId `liftM` GVec.basicUnsafeFreeze mv
basicUnsafeThaw (V_TermId v) = MV_TermId `liftM` GVec.basicUnsafeThaw v
basicLength (V_TermId v) = GVec.basicLength v
basicUnsafeSlice i l (V_TermId v) = V_TermId (GVec.basicUnsafeSlice i l v)
basicUnsafeIndexM (V_TermId v) i = TermId `liftM` GVec.basicUnsafeIndexM v i
basicUnsafeCopy (MV_TermId mv)
(V_TermId v) = GVec.basicUnsafeCopy mv v
elemseq (V_TermId v) x = GVec.elemseq v (unTermId x)
{-# INLINE basicUnsafeFreeze #-}
{-# INLINE basicUnsafeThaw #-}
{-# INLINE basicLength #-}
{-# INLINE basicUnsafeSlice #-}
{-# INLINE basicUnsafeIndexM #-}
{-# INLINE basicUnsafeCopy #-}
{-# INLINE elemseq #-}