content: {-
    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 TypeApplications #-}
  
  -- | Support for working with Thrift-generated predicate types
  module Glean.Typed.Predicate
    ( -- * class Predicate
      Predicate(..)
    , justId
    , predicateRef
    , predicateSourceType
      -- * Lookup PidOf
    , HasPredicates(..)
      -- ** exception
    , MissingPredicate(..), throwMissingPredicate
      -- * Vector Predicates
    , Predicates, SchemaPredicates, makePredicates
    ) where
  
  import Control.Concurrent
  import Control.Exception
  import Data.HashMap.Strict (HashMap)
  import qualified Data.HashMap.Strict as HashMap
  import Data.IORef
  import qualified Data.Map as Map
  import Data.Proxy
  import Data.Text (Text)
  import Data.Vector.Primitive (Vector)
  import qualified Data.Vector.Primitive as Vector
  import System.IO.Unsafe
  import TextShow
  
  import Glean.RTS.Types (invalidPid)
  import qualified Glean.Schema.Util as Angle
  import Glean.Types (PredicateRef(..), SchemaInfo(..))
  import Glean.Typed.Binary
  import Glean.Typed.Id
  import qualified Glean.Angle.Types as Angle
  
  -- -----------------------------------------------------------------------------
  
  -- | 'Predicate' provides information about the schema's data structures,
  -- and convenience functions.
  --
  -- * Types
  --     * @'KeyType' Foo@ : the generated type (often @Foo_key*) that
  --       stores the predicate's fields.
  --     * @'ValueType' Foo@ : the generated values type, defaults to @()@.
  -- * Methods
  --     * @'getName' Foo@ : the @glean.thrift@ derived 'PredicateRef' holding
  --       the Text name and Int version of the predicate.
  --     * @'getIndex' Foo@ : internal use, index into all schema predicate list.
  --       Used to efficiently lookup @'T.PidOf' p@ for a particular repo.
  --
  -- Superclass constraint of 'BT.Type' ensures this can be serialized into
  -- the underlying binary format.
  
  class (Type p, Type (KeyType p), Type (ValueType p)) => Predicate p where
    type KeyType p
    type ValueType p
    type ValueType p = ()
    getName :: proxy p -> PredicateRef
    mkFact :: IdOf p -> Maybe (KeyType p) -> Maybe (ValueType p) -> p
    getId :: p -> IdOf p
    getFactKey :: p -> Maybe (KeyType p)
    getFactValue :: p -> Maybe (ValueType p)
  
    -- A unique integer per predicate, used for fast lookup of Pids when
    -- writing facts.
    getIndex :: proxy p -> PredicateIndex
    getIndex = const (predicateIndex (getName (Proxy :: Proxy p)))
      -- this should be computed once and cached per predicate type
  
  -- | Retrieve the Angle representation of the type
  predicateSourceType :: Predicate p => Proxy p -> Angle.SourceType' ()
  predicateSourceType proxy =
    Angle.PredicateTy () (Angle.convertRef (getName proxy))
  
  justId :: Predicate p => IdOf p -> p
  justId x = mkFact x Nothing Nothing
  
  newtype MissingPredicate = MissingPredicate PredicateRef deriving(Show)
  instance Exception MissingPredicate
  
  -- | Useful error for 'getPid' in 'HasPredicates'
  throwMissingPredicate :: forall p. Predicate p => PidOf p
  throwMissingPredicate = throw $ MissingPredicate $ getName (Proxy :: Proxy p)
  
  -- | An index into the Predicates vector. Used when writing facts only.
  type PredicateIndex = Int
  
  newtype Predicates = Predicates (Vector Pid)
    -- invalidPid for predicates missing from the schema
  
  -- | Class of values 'a' that have definitions for any predicate 'p'
  class HasPredicates a where
    getPid :: Predicate p => a -> PidOf p
  
  instance HasPredicates Predicates where
    getPid = get
      where
        get :: forall p. Predicate p => Predicates -> PidOf p
        get (Predicates pids)
          | pid /= invalidPid = PidOf pid
          | otherwise = throw $ MissingPredicate $ getName (Proxy :: Proxy p)
          where
            index = getIndex (Proxy :: Proxy p)
            pid = if index < Vector.length pids
              then pids Vector.! index
              else invalidPid
  
  -- | The type of @allPredicates@ generated for each schema. To get this,
  -- import the generated module for the schema, e.g. @Glean.Schema.Src@.
  type SchemaPredicates = [PredicateRef]
  
  predicateIndex :: PredicateRef -> Int
  predicateIndex ref =
    unsafeDupablePerformIO $ do
      modifyMVar predicateIndices $ \hm ->
        case HashMap.lookup ref hm of
          Nothing -> do
            n <- atomicModifyIORef' predicateNextIndex (\x -> (x+1, x+1))
            return (HashMap.insert ref n hm, n)
          Just n -> return (hm, n)
  
  {-# NOINLINE predicateIndices #-}
  predicateIndices :: MVar (HashMap PredicateRef Int)
  predicateIndices = unsafePerformIO $ newMVar HashMap.empty
  
  {-# NOINLINE predicateNextIndex #-}
  predicateNextIndex :: IORef Int
  predicateNextIndex = unsafePerformIO $ newIORef 0
  
  makePredicates :: [SchemaPredicates] -> SchemaInfo -> Predicates
  makePredicates schemas info = Predicates $
    (Vector.//) (Vector.replicate (maxIx+1) invalidPid) assocs
    where
      maxIx = maximum $ map fst assocs
      assocs =
        [ (predicateIndex ref, HashMap.lookupDefault invalidPid ref ref_ids)
        | ref <- concat schemas
        ]
      ref_ids = HashMap.fromList
        [(ref, Pid id) | (id,ref) <- Map.toList $ schemaInfo_predicateIds info]
  
  -- | Convert a predicate type to a versioned predicate name. This is
  -- useful for building Angle queries that refer to specific versions.
  predicateRef :: forall p proxy . Predicate p => proxy p -> Text
  predicateRef _ = name <> "." <> showt version
    where !(PredicateRef name version) = getName (Proxy @p)
  