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.
  -}
  
  module Glean.Database.Schema.ComputeIds
    ( computeIds
    , RefToIdEnv(..)
    , emptyRefToIdEnv
    , ResolvedSchemaId
    , refsToIds
    , HashedSchema(..)
    , emptyHashedSchema
    , RefTargetId
    ) where
  
  import Control.Monad
  import Control.Monad.State as State
  import Data.Bifunctor
  import Data.Bifoldable
  import Data.Graph
  import Data.HashMap.Strict (HashMap)
  import qualified Data.HashMap.Strict as HashMap
  import Data.List
  import qualified Data.Map as Map
  import qualified Data.Set as Set
  import qualified Data.Text as Text
  
  import Glean.Angle.Hash
  import Glean.Angle.Types as Schema
  import Glean.Schema.Util (showRef)
  import Glean.Schema.Types
  import Glean.Types
  
  -- | The schema with all predicate/type references replaced with
  -- PredicateId/TypeId.
  data HashedSchema = HashedSchema
    { hashedTypes :: HashMap TypeId TypeDef
    , hashedPreds :: HashMap PredicateId PredicateDef
    , schemaRefToIdEnv :: RefToIdEnv
    , hashedSchemaEnv :: NameEnv RefTargetId
    , hashedSchemaAllVersion :: Version
    , hashedSchemaId :: SchemaId
    }
  
  type RefTargetId = RefTarget PredicateId TypeId
  
  emptyHashedSchema :: HashedSchema
  emptyHashedSchema = HashedSchema
    { hashedTypes = HashMap.empty
    , hashedPreds =  HashMap.empty
    , schemaRefToIdEnv = emptyRefToIdEnv
    , hashedSchemaEnv = HashMap.empty
    , hashedSchemaAllVersion = 0
    , hashedSchemaId = SchemaId ""
    }
  
  type Def_ p t = RefTarget p t
  
  type DefResolved = Def_ ResolvedPredicateDef ResolvedTypeDef
  type Def = Def_ PredicateDef TypeDef
  
  data RefToIdEnv = RefToIdEnv
    { typeRefToId :: HashMap TypeRef TypeId
    , predRefToId :: HashMap PredicateRef PredicateId
    }
  
  emptyRefToIdEnv :: RefToIdEnv
  emptyRefToIdEnv = RefToIdEnv HashMap.empty HashMap.empty
  
  type ResolvedSchemaId = ResolvedSchema PredicateId TypeId
  
  refsToIds
    :: Bifunctor f
    => RefToIdEnv
    -> f PredicateRef TypeRef
    -> f PredicateId TypeId
  refsToIds (RefToIdEnv tids pids) =
    bimap (lookupPredicateId pids) (lookupTypeId tids)
    where
      lookupPredicateId m p = HashMap.findWithDefault err p m
        where
          err = error $ "lookupPredicateId: " <> Text.unpack (showRef p)
      lookupTypeId m p = HashMap.findWithDefault err p m
        where
          err = error $ "lookupTypeId: " <> Text.unpack (showRef p)
  
  attachDerivations
    :: [ResolvedSchemaRef]
    -> [HashMap PredicateRef ResolvedPredicateDef]
  attachDerivations schemas =
    map (HashMap.mapWithKey attach . resolvedSchemaPredicates) schemas
    where
    allDerivings = HashMap.unions (map resolvedSchemaDeriving schemas)
  
    attach ref def = case HashMap.lookup ref allDerivings of
      Just drv | not (isDefaultDeriving $ derivingDefDeriveInfo drv) ->
        def { predicateDefDeriving =  derivingDefDeriveInfo drv }
        -- see Note [overriding default deriving]
      _ -> def
  
  isDefaultDeriving :: DerivingInfo q -> Bool
  isDefaultDeriving (Derive DeriveIfEmpty _) = True
  isDefaultDeriving _ = False
  
  -- | Compute the PredicateId / TypeId for each definition, and substitute for
  -- PredicateRef/TypeRef with PredicateId/TypeId inside all the definitions
  
  computeIds
    :: (forall k v . HashMap k v -> [(k,v)])
    -> [ResolvedSchemaRef]
    -> Maybe (SchemaId, Version)
    -> HashedSchema
  computeIds toList schemas versions = flip evalState emptyRefToIdEnv $ do
    let
      preds = attachDerivations schemas
  
      -- first topologically sort the predicates and types, so that we can
      -- process them in dependency order.
      edges :: [(DefResolved, RefResolved, [RefResolved])]
      edges =
        [ (RefPred def, RefPred (predicateDefRef def), predicateDefRefs def)
        | m <- preds
        , (_, def) <- toList m
        ] ++
        [ (RefType def, RefType (typeDefRef def), typeDefRefs def)
        | m <- map resolvedSchemaTypes schemas
        , (_, def) <- toList m
        ]
  
      collectRefs = bifoldMap ((:[]) . RefPred) ((:[]) . RefType)
  
      predicateDefRefs def =
        collectRefs (predicateDefKeyType def) <>
        collectRefs (predicateDefValueType def) <>
        foldMap collectRefs (predicateDefDeriving def)
  
      typeDefRefs def = collectRefs (typeDefType def)
  
    defs <- fmap concat $ forM (stronglyConnComp edges) $ \comp ->
      case comp of
        AcyclicSCC def -> do
          resolved <- resolveDef def
          let newdef = updateDefWithHash (fingerprintDef resolved) resolved
          State.modify (extend newdef)
          return [newdef]
        CyclicSCC defs -> do
          -- sort the definitions in the cycle by ref, so that the next
          -- steps produce deterministic hashes.
          let sorted = sortOn (bimap predicateDefRef typeDefRef) defs
          -- first map every ref in the cycle to the Id hash0
          defs1 <-
            withState (extends (zip defs (repeat hash0))) $
              mapM resolveDef sorted
          -- Compute a hash for each def
          let hashes = map fingerprintDef defs1
          -- Next make a hash of the whole cycle.
          let cycleHash = hashBinary hashes
          -- Note that it's possible to have two predicates with the
          -- same name (different versions) in the cycle, and we
          -- definitely want them to end up with different hashes
          -- (unless they have identical representations) so the hash of
          -- each declaration in the cycle is the hash of the
          -- declaration plus the hash of the whole cycle.
          let hashes2 =
                [ (def, hashBinary (hash1, cycleHash))
                | (def, hash1) <- zip sorted hashes ]
          State.modify (extends hashes2)
          -- now resolve the defs again with the correct Ids
          forM hashes2 $ \(def, hash) ->
            updateDefWithHash hash <$> resolveDef def
  
    env <- State.get
  
    let preds = HashMap.fromList
          [ (predicateDefRef def, def) | RefPred def <- defs]
  
        -- see Note [overriding default deriving]
        attachDefaultDerivings preds = foldr attach preds
          [ (id, derivingDefDeriveInfo drv)
          | schema <- schemas
          , (ref, drv) <- toList (resolvedSchemaDeriving schema)
          , isDefaultDeriving $ derivingDefDeriveInfo drv
          , Just id <- [HashMap.lookup ref (predRefToId env)]
          ]
          where
          attach (id,drv) = HashMap.adjust f id
            where
            f (PredicateDef id key val _ s) =
              PredicateDef id key val (fmap (refsToIds env) drv) s
  
        (allVersion, schemaId, schemaEnv) = makeSchemaEnv schemas versions env
  
    return HashedSchema {
        hashedTypes = HashMap.fromList
          [ (typeDefRef def, def) | RefType def <- defs],
        hashedPreds = attachDefaultDerivings preds,
        schemaRefToIdEnv = env,
        hashedSchemaEnv = schemaEnv,
        hashedSchemaAllVersion = allVersion,
        hashedSchemaId = schemaId
      }
  
    where
      updateDefWithHash :: Hash -> Def -> Def
      updateDefWithHash hash def = case def of
          RefType (TypeDef (TypeId ref _) ty s) ->
            RefType (TypeDef (TypeId ref hash) ty s)
          RefPred (PredicateDef (PredicateId ref _) key val drv s) ->
            RefPred (PredicateDef (PredicateId ref hash) key val drv s)
  
      extend :: Def -> RefToIdEnv -> RefToIdEnv
      extend def (RefToIdEnv tids pids) =
        case def of
          RefType (TypeDef id _ _) ->
            RefToIdEnv (HashMap.insert (typeIdRef id) id tids) pids
          RefPred (PredicateDef id _ _ _ _) ->
            RefToIdEnv tids (HashMap.insert (predicateIdRef id) id pids)
  
      extends :: [(DefResolved, Hash)] -> RefToIdEnv -> RefToIdEnv
      extends list env = foldr one env list
        where
        one (def, hash) (RefToIdEnv tids pids) =
          case def of
            RefType (TypeDef ref _ _) ->
              let id = TypeId ref hash in
              RefToIdEnv (HashMap.insert ref id tids) pids
            RefPred (PredicateDef ref _ _ _ _) ->
              let id = PredicateId ref hash in
              RefToIdEnv tids (HashMap.insert ref id pids)
  
      resolveDef :: DefResolved -> State RefToIdEnv Def
      resolveDef (RefType (TypeDef ref ty s)) = do
        env <- State.get
        let
          newTy = refsToIds env ty
          newId = TypeId {
            typeIdRef = ref,
            typeIdHash = hash0 }
          def = RefType (TypeDef newId newTy s)
        return def
      resolveDef (RefPred (PredicateDef ref key val drv s)) = do
        env <- State.get
        let
          newKey = refsToIds env key
          newVal = refsToIds env val
          newDrv = fmap (resolveQuery env) drv
          newId = PredicateId {
            predicateIdRef = ref,
            predicateIdHash = hash0 }
          def = RefPred (PredicateDef newId newKey newVal newDrv s)
        return def
  
  resolveQuery
    :: RefToIdEnv
    -> Query_ PredicateRef TypeRef
    -> Query_ PredicateId TypeId
  resolveQuery env (SourceQuery head stmts ord) =
    SourceQuery (refsToIds env <$> head) (refsToIds env <$> stmts) ord
  
  -- Serialize a definition to produce its hash.  Note that the hash
  -- includes the PredicateRef/TypeRef: two predicates or types are the
  -- same only if they the same name, same version, and same
  -- representation.
  fingerprintDef :: Def -> Hash
  fingerprintDef (RefType (TypeDef id ty _)) =
    hashBinary (showRef ref, rmLocType ty)
    where ref = typeIdRef id
  fingerprintDef (RefPred (PredicateDef id keyTy valTy drv _)) =
    hashBinary (showRef ref,rmLocType keyTy, rmLocType valTy, fmap rmLocQuery drv)
    where ref = predicateIdRef id
  
  {- Note [overriding default deriving]
  
  If we add a "deriving P default" declaration for a predicate P to the
  schema, after a DB was created with P, we would like this to work: the
  new P should get the same Pid, so queries still work.
  
  But, if the derivation is included in the hash, the new P with the
  deriving decl will get a different hash. It's difficult to resolve
  this later, because we only want the derivation to take effect if the
  two Ps are really the same (they have the same types).
  
  So the hack to solve this for now is to ignore the "derive default"
  derivations when computing hashes, and add them back in
  afterwards. The old and the new Ps will get the same hashes, if they
  really only differ in their derivations.
  
  -}
  
  
  makeSchemaEnv
    :: [ResolvedSchemaRef]
    -> Maybe (SchemaId, Version)
    -> RefToIdEnv
    -> (Version, SchemaId, NameEnv RefTargetId)
  makeSchemaEnv resolved versions refToIdEnv
    | Just (schemaId, ver) <- versions,
      Just schema <- Map.lookup ver resolvedAlls =
      (ver, schemaId, mkEnv schema)
    | Just (ver, schema) <- Map.lookupMax resolvedAlls =
      let
        env = mkEnv schema
        schemaId = SchemaId $ Text.pack $ show $ hashNameEnv env
      in
      (fromIntegral ver, schemaId, env)
    | otherwise =
      error "no \"all\" schema"
    where
      mkEnv ResolvedSchema{..} =
         HashMap.union
           (mapNameEnv (Just . refsToIds refToIdEnv)
             resolvedSchemaQualScope)
           versionedNameEnv
  
      resolvedAlls = Map.fromList
        [ (resolvedSchemaVersion schema, schema)
        | schema <- resolved
        , resolvedSchemaName schema == "all" ]
  
      -- In the environment that we use for resolving names in queries
      -- and derivations later, we need to accept explicitly versioned
      -- references to predicates and types even for those predicates
      -- and types that are not visible in the scope of the "all" schema.
      versionedNameEnv = HashMap.fromList
        [ (SourceRef name (Just ver), Set.singleton target)
        | (name, ver, target) <-
            [ (predicateRef_name r, predicateRef_version r, RefPred t)
            | (r, t) <- HashMap.toList (predRefToId refToIdEnv) ] ++
            [ (typeRef_name r, typeRef_version r, RefType t)
            | (r, t) <- HashMap.toList (typeRefToId refToIdEnv) ]
        ]
  
  -- | How to take a NameEnv and compute the schema hash from it. This
  -- hash uniquely identifies a particular NameEnv that can be used to
  -- resolve a type or predicate name, or in general a query.
  hashNameEnv :: NameEnv (RefTarget PredicateId TypeId) -> Hash
  hashNameEnv env =
    hashBinary (sort (HashMap.toList env))
  