{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# LANGUAGE DeriveTraversable, GeneralizedNewtypeDeriving #-}
-- | A tagged intmap to keep the type of keys clear.  Mainly intended for
-- Glean predicates.
module Glean.Util.PredMap
  ( PredMap(..), makePredMap, fromList, fromDistinctAscList, fromListWith
  , lookup, findWithDefault, mapMaybe, alter, empty, insert, insertWith, elems
  , union
  , unionsWith
  , keys
  , toList, toAscList
  , size
  , member, notMember
  , filterValues
  ) where

import Control.DeepSeq (NFData)
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import qualified Data.Maybe (mapMaybe)
import Prelude hiding (lookup)

import Glean hiding (keys)

-- -----------------------------------------------------------------------------
-- Helpers, to make most of this module into one liners

one :: IdOf p -> Int
one = fromIntegral . fromFid . idOf

withKey :: (IdOf p, v) -> (Int, v)
withKey (p, v) = (one p, v)

-- -----------------------------------------------------------------------------

-- | I want some more type safety around IntMap. The 'p' will usually
-- be a Predicate type, e.g. (IdOf p) as the logical key.
newtype PredMap p v = PredMap { predMap :: IntMap v }
  deriving (Show, Semigroup, Monoid, Functor, Foldable, Traversable, NFData)

empty :: PredMap p v
empty = PredMap IntMap.empty

-- | Simplest way to make a 'PredMap' but not the only way.  Any 'p' without
-- a key are silently dropped.
makePredMap :: Predicate p => [p] -> PredMap p (KeyType p)
makePredMap = PredMap . IntMap.fromList . Data.Maybe.mapMaybe getWithKey
  where getWithKey p = (one (getId p),) <$> getFactKey p

fromList :: [(IdOf p, v)] -> PredMap p v
fromList = PredMap . IntMap.fromList . map withKey

fromDistinctAscList :: [(IdOf p, v)] -> PredMap p v
fromDistinctAscList = PredMap . IntMap.fromDistinctAscList . map withKey

fromListWith :: (v -> v -> v) -> [(IdOf p, v)] -> PredMap p v
fromListWith f = PredMap . IntMap.fromListWith f . map withKey

lookup :: Predicate p => IdOf p -> PredMap p v -> Maybe v
lookup p pm = IntMap.lookup (one p) (predMap pm)

findWithDefault :: Predicate p => v -> IdOf p -> PredMap p v -> v
findWithDefault d p = IntMap.findWithDefault d (one p) . predMap

mapMaybe :: Predicate p => (a -> Maybe b) -> PredMap p a -> PredMap p b
mapMaybe f = PredMap . IntMap.mapMaybe f . predMap

alter :: (Maybe v -> Maybe v) -> IdOf p -> PredMap p v -> PredMap p v
alter f p = PredMap . IntMap.alter f (one p) . predMap

insert :: IdOf p -> v -> PredMap p v -> PredMap p v
insert p v = PredMap . IntMap.insert (one p) v . predMap

insertWith :: (v -> v -> v) -> IdOf p -> v -> PredMap p v -> PredMap p v
insertWith f p v = PredMap . IntMap.insertWith f (one p) v . predMap

union :: PredMap p v -> PredMap p v -> PredMap p v
union a b = PredMap $ IntMap.union (predMap a) (predMap b)

unionsWith :: (v -> v -> v) -> [PredMap p v] -> PredMap p v
unionsWith f = PredMap . IntMap.unionsWith f . map predMap

keys :: PredMap p v -> [IdOf p]
keys = map (\k -> IdOf (Fid (fromIntegral k))) . IntMap.keys . predMap

elems :: PredMap p v -> [v]
elems = IntMap.elems . predMap

-- | Synonym of 'toAscList'
toList :: PredMap p v -> [(IdOf p, v)]
toList = toAscList

toAscList :: PredMap p v -> [(IdOf p, v)]
toAscList = map (\ (k, v) -> (IdOf (Fid (fromIntegral k)), v))
       . IntMap.toAscList . predMap

size :: PredMap p v -> Int
size = IntMap.size . predMap

member :: IdOf p -> PredMap p v -> Bool
member p pm = IntMap.member (one p) (predMap pm)

notMember :: IdOf p -> PredMap p v -> Bool
notMember p pm = IntMap.notMember (one p) (predMap pm)

filterValues :: (v -> Bool) -> PredMap p v -> PredMap p v
filterValues f (PredMap m) = PredMap $ IntMap.filter f m