InvertedIndex.hs

module NixTree.Data.InvertedIndex
  ( InvertedIndex,
    iiFromList,
    iiInsert,
    iiSearch,
  )
where

import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Text as Text

data InvertedIndex a = InvertedIndex
  { iiElems :: Map Text a,
    iiUnigrams :: Map Char (Set Text),
    iiBigrams :: Map (Char, Char) (Set Text),
    iiTrigrams :: Map (Char, Char, Char) (Set Text)
  }
  deriving (Generic, Show)

instance (NFData a) => NFData (InvertedIndex a)

iiInsert :: Text -> a -> InvertedIndex a -> InvertedIndex a
iiInsert txt val InvertedIndex {iiElems, iiUnigrams, iiBigrams, iiTrigrams} =
  InvertedIndex
    { iiElems = Map.insert txt val iiElems,
      iiUnigrams = combine iiUnigrams (unigramsOf txt),
      iiBigrams = combine iiBigrams (bigramsOf txt),
      iiTrigrams = combine iiTrigrams (trigramsOf txt)
    }
  where
    combine orig chrs =
      Map.unionWith
        (<>)
        orig
        (setToMap (Set.singleton txt) chrs)

iiFromList :: (Foldable f) => f (Text, a) -> InvertedIndex a
iiFromList =
  foldl'
    (flip (uncurry iiInsert))
    (InvertedIndex Map.empty Map.empty Map.empty Map.empty)

setToMap :: v -> Set k -> Map k v
setToMap v = Map.fromDistinctAscList . map (,v) . Set.toAscList

unigramsOf :: Text -> Set Char
unigramsOf = Set.fromList . Text.unpack . Text.toLower

bigramsOf :: Text -> Set (Char, Char)
bigramsOf txt = case Text.unpack (Text.toLower txt) of
  p1@(_ : p2) -> Set.fromList $ zip p1 p2
  _ -> Set.empty

trigramsOf :: Text -> Set (Char, Char, Char)
trigramsOf txt = case Text.unpack (Text.toLower txt) of
  p1@(_ : p2@(_ : p3)) -> Set.fromList $ zip3 p1 p2 p3
  _ -> Set.empty

iiSearch :: forall a. Text -> InvertedIndex a -> Map Text a
iiSearch txt InvertedIndex {iiElems, iiUnigrams, iiBigrams, iiTrigrams}
  | Text.length txt == 0 = iiElems
  | Text.length txt == 1 = using unigramsOf iiUnigrams
  | Text.length txt == 2 = using bigramsOf iiBigrams
  | otherwise = using trigramsOf iiTrigrams
  where
    lowerTxt = Text.toLower txt
    using :: (Ord c) => (Text -> Set c) -> Map c (Set Text) -> Map Text a
    using getGrams m =
      Map.intersection m (setToMap () (getGrams txt))
        & Map.elems
        & \case
          [] -> Set.empty
          x : xs -> foldl' Set.intersection x xs
        & Set.filter (\t -> lowerTxt `Text.isInfixOf` Text.toLower t)
        & Map.restrictKeys iiElems