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 ApplicativeDo #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE DeriveGeneric #-}
  
  module Glean.Glass.SearchRelated
    ( searchRelatedEntities
    , searchInheritedEntities
    , searchRecursiveEntities
    , Recursive(..)
    , SearchStyle(..)
    , RelatedLocatedEntities(..)
    , RelatedEntities(..)
    , LocatedEntity
    , InheritedContainer
    , edgesToTopoMap
    ) where
  
  import Control.Monad (forM)
  import Control.Monad.Catch (MonadThrow(throwM))
  import Data.Hashable ( Hashable(..) )
  import Data.HashSet ( HashSet )
  import Data.List (groupBy, sortOn)
  import Data.Function (on)
  import Data.HashMap.Strict ( HashMap )
  import qualified Data.HashMap.Strict as HM
  import GHC.Generics (Generic)
  import qualified Data.HashSet as HashSet
  
  import qualified Glean
  import Glean.Angle as Angle
  import qualified Glean.Schema.CodemarkupTypes.Types as Code
  import qualified Glean.Schema.Codemarkup.Types as Code
  import qualified Glean.Schema.Code.Types as Code
  
  import Glean.Glass.Search.Class
  import Glean.Glass.Base (GleanPath (..))
  import Glean.Glass.Path
  import Glean.Glass.SymbolId (entityToAngle, toSymbolId)
  import Glean.Glass.Types
  import Glean.Glass.Utils (searchRecursiveWithLimit)
  import Glean.Haxl.Repos (RepoHaxl)
  
  -- | Whether to expand relationships recursively
  data Recursive
    = Recursive
    | NotRecursive
    deriving Eq
  
  -- | Whether to hide uninteresting things
  data SearchStyle
    = ShowAll
    | HideUninteresting
  
  -- | Pairs of edges of related entities
  data RelatedEntities = RelatedEntities
    { parentEntity :: !Code.Entity
    , parentLocation :: {-# UNPACK #-}!Code.Location
    , childEntity :: !Code.Entity
    , childLocation :: {-# UNPACK #-}!Code.Location
    } deriving (Eq,Generic,Hashable)
  
  -- | Pairs of edges of related entities, with all metadata
  data RelatedLocatedEntities = RelatedLocatedEntities
    { parentRL :: LocatedEntity
    , childRL :: LocatedEntity
    } deriving (Generic,Eq,Hashable)
  
  type LocatedEntity = (ResultLocation Code.Entity, SymbolId)
  
  -- Convenience type for a parent with a set of contained children
  type InheritedContainer = (LocatedEntity, [LocatedEntity])
  
  styleToCode :: SearchStyle -> Code.SearchStyle
  styleToCode ShowAll = Code.SearchStyle_ShowAll
  styleToCode HideUninteresting = Code.SearchStyle_HideUninteresting
  
  -- | Flatten results into the container-level topological ordering
  -- from child to set of parents it inherits from.
  edgesToTopoMap
    :: [RelatedLocatedEntities] -> HashMap SymbolId (HashSet SymbolId)
  edgesToTopoMap edges = HM.fromListWith HashSet.union
    [ (snd (childRL e), HashSet.singleton (snd (parentRL e))) | e <- edges ]
  
  -- returns all inherited members in container entity as "RelatedEntity"
  -- parentEntity is the entity defining the member
  -- childEntity is the member itself
  runSearchInherited
    :: SearchStyle
    -> Int
    -> Angle Code.Entity
    -> RepoHaxl u w [RelatedEntities]
  runSearchInherited style limit angle = do
    (entities, _truncated) <- searchRecursiveWithLimit (Just limit) $
      query angle
    pure $
      [ RelatedEntities
        { parentEntity = parentEntity_parent
        , childEntity = childEntity_child
        , parentLocation = parentEntity_location
        , childLocation = childEntity_location
        }
      | (Code.ParentEntity{..}, Code.ChildEntity{..}) <- entities ]
    where
      styleTy = styleToCode style
      query
        :: Angle Code.Entity
        -> Angle (Code.ParentEntity, Code.ChildEntity)
      query entity =
        vars $ \(parent :: Angle Code.ParentEntity)
            (child :: Angle Code.ChildEntity) ->
          tuple (parent, child) `where_` [
            wild .= predicate @Code.SearchInheritedEntities (
              rec $
                field @"base" entity $
                field @"parent" parent $
                field @"child" child $
                field @"style" (enum styleTy)
              end)
            ]
  
  -- returns all inherited members in container entity
  -- grouped by parentEntity
  searchInheritedEntities
    :: SearchStyle
    -> Int
    -> Code.Entity
    -> RepoName
    -> RepoHaxl u w [InheritedContainer]
  searchInheritedEntities style limit entity repo = do
    angle <- case entityToAngle entity of
          Right angle -> return angle
          Left t -> throwM (ServerException t)
    relatedEntities <- runSearchInherited style limit angle >>= toSymbolIds repo
    return $ groupChildEntities relatedEntities
    where
      groupChildEntities :: [RelatedLocatedEntities] -> [InheritedContainer]
      groupChildEntities xs =
        map (\g -> (parentRL (head g), map childRL g)) grouped
        where
          sorted = sortOn parentRL xs
          grouped = groupBy ((==) `on` parentRL) sorted
  
  --
  -- Given some search parameters, find entities by relation
  --
  searchRelatedEntities
    :: Int
    -> SearchStyle
    -> Recursive
    -> RelationDirection
    -> RelationType
    -> Code.Entity
    -> RepoName
    -> RepoHaxl u w [RelatedLocatedEntities]
  searchRelatedEntities limit hide recurse dir rel entity repo =
    toSymbolIds repo =<< searchRelation opts limit [entity] HashSet.empty
    where
      opts = SearchOptions limit hide recurse rel dir
  
  -- | For internal searches, we don't need the symbol id. So we can be slightly
  -- more efficient. TODO: we could also avoid the location in the result entirely
  -- used for lucky search/context resolution
  searchRecursiveEntities
    :: Int -> RelationDirection -> RelationType -> Code.Entity
    -> RepoHaxl u w [RelatedEntities]
  searchRecursiveEntities limit dir rel entity =
      searchRelation opts limit [entity] HashSet.empty
    where
      opts = SearchOptions limit ShowAll Recursive rel dir
  
  -- | Lift entity search results into pairs of entities that we found,
  -- along with their location and symbol id
  --
  -- TODO: we will be recomputing these for some symbols many times,
  -- and in some cases we are dropping them as well. Can we defer this until
  -- later?
  --
  toSymbolIds
    :: RepoName -> [RelatedEntities] -> RepoHaxl u w [RelatedLocatedEntities]
  toSymbolIds repo edges = mapM locatePairs edges
    where
      locatePairs RelatedEntities{..} = RelatedLocatedEntities
        <$> mkLocate parentEntity parentLocation
        <*> mkLocate childEntity childLocation
      mkLocate entity Code.Location{..} = do
        symId <- symbol entity location_file
        return ((entity,location_file,location_location,location_name), symId)
      symbol entity file = do
        path <- GleanPath <$> Glean.keyOf file
        toSymbolId (fromGleanPath repo path) entity
  
  -- | Constants to determine how we search
  data SearchOptions = SearchOptions {
    totalLimit :: Int,
    style :: SearchStyle,
    recursive :: Recursive,
    relation :: RelationType,
    direction :: RelationDirection
  }
  
  --
  -- Search driver, expand search until done, returning pairs of edges
  -- of entity relationships.
  --
  searchRelation
    :: SearchOptions
    -> Int
    -> [Code.Entity]
    -> HashSet RelatedEntities
    -> RepoHaxl u w [RelatedEntities]
  searchRelation opts@SearchOptions{..} limit toVisit visited =
      runSearch searchFn opts limit toVisit visited
    where
      -- the search flavor is known up front, so we can partially apply
      baseSearchFn = \ty angle -> runSearchRelated totalLimit style angle ty
      searchFn = case (relation, direction) of
        (RelationType_Extends, RelationDirection_Parent) ->
          baseSearchFn Code.RelationType_ExtendsParentOfChild
        (RelationType_Extends, RelationDirection_Child) ->
          baseSearchFn Code.RelationType_ExtendsChildOfParent
        (RelationType_Contains, RelationDirection_Parent) ->
          baseSearchFn Code.RelationType_ContainsParentOfChild
        (RelationType_Contains, RelationDirection_Child) ->
          baseSearchFn Code.RelationType_ContainsChildOfParent
        (RelationType_RequireExtends, RelationDirection_Parent) ->
          baseSearchFn Code.RelationType_RequireExtendsParentOfChild
        (RelationType_RequireImplements, RelationDirection_Parent) ->
          baseSearchFn Code.RelationType_RequireImplementsParentOfChild
        (RelationType_RequireClass, RelationDirection_Parent) ->
          baseSearchFn Code.RelationType_RequireClassParentOfChild
        _ -> \_angle -> pure []
  
  -- recursively search up to limit
  runSearch
    :: ([Angle Code.Entity] -> RepoHaxl u w [RelatedEntities])
    -> SearchOptions
    -> Int
    -> [Code.Entity]
    -> HashSet RelatedEntities
    -> RepoHaxl u w [RelatedEntities]
  runSearch searchFn opts@SearchOptions{..} !limit toVisit visited = do
    angle <- forM toVisit $ \entity ->
      case entityToAngle entity of
        Right angle -> return angle
        Left t -> throwM (ServerException t)
    justVisited <- searchFn angle
    let newlyVisited = HashSet.fromList justVisited `HashSet.difference` visited
        visited' = visited `HashSet.union` newlyVisited
        toVisit = HashSet.toList $ case direction of
          RelationDirection_Parent -> HashSet.map parentEntity newlyVisited
          RelationDirection_Child -> HashSet.map childEntity newlyVisited
          _ -> HashSet.empty -- unknown direction
        recLimit = limit - length visited'
    if
      recursive == Recursive && recLimit > 0 && recLimit < limit &&
      not (null toVisit)
    then
      runSearch searchFn opts recLimit toVisit visited'
    else
      return $! HashSet.toList visited'
  
  runSearchRelated
    :: Int
    -> SearchStyle
    -> [Angle Code.Entity]
    -> Code.RelationType
    -> RepoHaxl u w [RelatedEntities]
  runSearchRelated limit style angle searchType = do
    (entities, _truncated) <- searchRecursiveWithLimit (Just limit) $
      searchRelatedEntitiesQ searchType styleType entities
    pure $
      [ RelatedEntities
        { parentEntity = parentEntity_parent
        , childEntity = childEntity_child
        , parentLocation = parentEntity_location
        , childLocation = childEntity_location
        }
      | Code.SearchRelatedEntities{..} <- entities
      , Just Code.SearchRelatedEntities_key{..} <- [searchRelatedEntities_key]
      , let Code.ParentEntity{..} = searchRelatedEntities_key_parent
      , let Code.ChildEntity{..} = searchRelatedEntities_key_child
      ]
    where
      entities = elementsOf (array angle)
      styleType = styleToCode style
  
  --
  -- unified search by relation
  --
  searchRelatedEntitiesQ
    :: Code.RelationType
    -> Code.SearchStyle
    -> Angle Code.Entity
    -> Angle Code.SearchRelatedEntities
  searchRelatedEntitiesQ queryTy styleTy entity = case queryTy of
      Code.RelationType_ExtendsParentOfChild -> keyedByChild
      Code.RelationType_ContainsParentOfChild -> keyedByChild
      Code.RelationType_ExtendsChildOfParent -> keyedByParent
      Code.RelationType_ContainsChildOfParent -> keyedByParent
      Code.RelationType_RequireClassParentOfChild -> keyedByChild
      Code.RelationType_RequireImplementsParentOfChild -> keyedByChild
      Code.RelationType_RequireExtendsParentOfChild -> keyedByChild
      Code.RelationType__UNKNOWN _ -> error "Unkonwn Code.RelationType"
    where
      keyedByChild = predicate @Code.SearchRelatedEntities $
        rec $ -- bind to child entity
          field @"query" (enum queryTy) $
          field @"style" (enum styleTy) $
          field @"child" (rec $ field @"child" entity end)
        end
      keyedByParent = predicate @Code.SearchRelatedEntities $
        rec $ -- search by parent entities
          field @"query" (enum queryTy) $
          field @"style" (enum styleTy) $
          field @"parent" (rec $ field @"parent" entity end)
        end
  