{- 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 QuasiQuotes #-} {-# LANGUAGE TypeApplications #-} module Angle.OptTest (main) where import Data.Default import Data.Int import Data.Map (Map) import qualified Data.Map as Map import Data.Maybe import Data.Text (Text) import Test.HUnit import TestRunner import Util.String.Quasi import Glean.Backend.Types as Backend import Glean.Init import Glean.Query.Angle as Angle import Glean.Query.Thrift as Thrift import Glean.Query.Thrift.Internal import qualified Glean.Schema.Cxx1.Types as Cxx import qualified Glean.Schema.GleanTest.Types as Glean.Test import Glean.Typed hiding (end) import Glean.Types import TestDB main :: IO () main = withUnitTest $ testRunner $ TestList [ TestLabel "opt" optTest ] queryStats :: forall q backend . (Backend backend) => backend -> Repo -> Query q -> IO ([q], Maybe UserQueryStats) queryStats be repo (Query query) = do let opts = fromMaybe def (userQuery_options query) query' = query { userQuery_encodings = [UserQueryEncoding_bin def] , userQuery_options = Just opts { userQueryOptions_collect_facts_searched = True } } UserQueryResults{..} <- userQuery be repo query' results <- decodeResults userQueryResults_results decodeAsFact return (results, userQueryResults_stats) factsSearched :: PredicateRef -> Map PredicateRef Int64 -> Maybe UserQueryStats -> Maybe Int factsSearched ref lookupPid maybeStats = do pid <- Map.lookup ref lookupPid stats <- maybeStats searched <- userQueryStats_facts_searched stats count <- Map.lookup pid searched return (fromIntegral count) optTest :: Test optTest = dbTestCase $ \env repo -> do si <- getSchemaInfo env (Just repo) def { getSchemaInfo_omit_source = True } let lookupPid = Map.fromList [ (ref,pid) | (pid,ref) <- Map.toList (schemaInfo_predicateIds si) ] -- Test unification: this query should be as efficient as simply -- cxx1.Name "an".. (_, stats) <- queryStats env repo $ angle @Cxx.Name [s| N where {Y,_} = W : {a : string, b : nat}; N = cxx1.Name X; X = Y; W = {"an".., 3} |] assertEqual "opt 1" (Just 2) $ factsSearched (PredicateRef "cxx1.Name" 5) lookupPid stats -- test that unification can optimise A = B | C, unifying A/B and A/C (_, stats) <- queryStats env repo $ angle @Cxx.Name [s| N where {"an".., N} = {X:string, cxx1.Name X} | {"def", cxx1.Name "def"} |] assertEqual "opt 2" (Just 2) $ factsSearched (PredicateRef "cxx1.Name" 5) lookupPid stats -- test for recursive substitution r <- runQuery_ env repo $ angleData @Nat $ [s| X where X = Y : nat; Y = X; X = 3 |] assertEqual "opt 3" [Nat 3] r -- test for optimised ordering of nested fact patterns (_, stats) <- queryStats env repo $ angleData @Text [s| S where glean.test.Predicate { string_ = "clobber", sum_ = { c = { string_ = S }} } |] assertEqual "opt 4" (Just 4) $ factsSearched (PredicateRef "glean.test.Predicate" 6) lookupPid stats -- test for optimising away unmatchable alternative without -- removing matchable ones r <- runQuery_ env repo $ angleData @Nat [s| 1 = (2|1); 1 |] assertEqual "opt 5" 1 (length r) -- Test for optimising away unmatchable alternatives. (_, stats) <- queryStats env repo $ angleData @Text [s| B where { "a", B } = # this should be optimised away: ( { A, B } where glean.test.StringPair { A, B }; A = "b") | # leaving this only: ( { A, B } where glean.test.StringPair { A, B }; A = "a") |] -- we can tell the first alternative was optimised away by counting -- the number of facts searched for glean.test.StringPair. assertEqual "optimise or 1" (Just 1) $ factsSearched (PredicateRef "glean.test.StringPair" 6) lookupPid stats -- Test that unification works properly with literal fact IDs result : _ <- runQuery_ env repo (allFacts :: Query Glean.Test.StringPair) let fid = factId (getId result) results <- runQuery_ env repo $ Angle.query $ fid `where_` [ fid .= fid ] assertEqual "unify fact Id" 1 (length results) -- Test that trivially false negations fail the entire query (_, stats) <- queryStats env repo $ angle @Glean.Test.StringPair [s| A where A = glean.test.StringPair _; !(1 = 1); |] assertEqual "opt - negation" Nothing $ factsSearched (PredicateRef "glean.test.StringPair" 6) lookupPid stats -- Test for unification inside negation (_, stats) <- queryStats env repo $ angleData @Nat [s| A where !(!(glean.test.Node { label = X }; X = "g"..)); A = 1; |] assertEqual "opt - negation 2" (Just 1) $ factsSearched (PredicateRef "glean.test.Node" 6) lookupPid stats -- a false condition in an if pattern gets optimised away. (_, stats) <- queryStats env repo $ angleData @Nat [s| if (1 = 2; glean.test.StringPair _) then 1 else 1 |] assertEqual "optimise if condition" Nothing $ factsSearched (PredicateRef "glean.test.StringPair" 6) lookupPid stats -- we should recognise that the `cxx1.Name _` statement only checks -- whether there is any cxx1.Name fact in the db and stop after finding -- the first result. (_, stats) <- queryStats env repo $ angle @Glean.Test.Predicate [s| cxx1.Name _; glean.test.Predicate _ |] assertEqual "optimise existence check" (Just 1) $ factsSearched (PredicateRef "cxx1.Name" 5) lookupPid stats -- check that we unify multiple fact generators (_, stats) <- queryStats env repo $ angleData @(Text, Text) [s| {A, B} where X = glean.test.StringPair { _, "b" }; X = glean.test.StringPair { "a", _ }; X = glean.test.StringPair { A, B }; |] assertEqual "optimise multiple generators" (Just 1) $ factsSearched (PredicateRef "glean.test.StringPair" 6) lookupPid stats -- Check that the optimiser doesn't get too excited and eliminate generators -- that don't bind any variables. These can still fail if they don't match -- any facts. results <- runQuery_ env repo $ angleData @Nat [s| 1 where glean.test.StringPair { "none", _ } |] assertEqual "check that we don't eliminate empty fact generators" 0 (length results)