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 QuasiQuotes #-}
  {-# LANGUAGE TypeApplications #-}
  module Angle.DerivedTest (main) where
  
  import Data.List
  import Data.Text (Text)
  import Test.HUnit
  
  import TestRunner
  import Util.String.Quasi
  
  import Glean.Init
  import Glean.Query.Thrift as Thrift
  import qualified Glean.Schema.GleanTest.Types as Glean.Test
  import Glean.Typed hiding (end)
  
  import TestDB
  
  main :: IO ()
  main = withUnitTest $ testRunner $ TestList
    [ TestLabel "derived" $ derivedTest id
    , TestLabel "derived/page" $ derivedTest (limit 1)
    ]
  
  -- | Tests for derived predicates using Angle syntax. See
  -- JSONQueryTest.jsonDerivedTest for the Thrift-query-type versions of
  -- these tests.
  derivedTest :: (forall a . Query a -> Query a) -> Test
  derivedTest modify = dbTestCase $ \env repo -> do
    -- Several tests with the trivial derived predicate: IsThree ::= 3
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.IsThree
      [s|
        glean.test.IsThree _
      |]
    print results
    assertEqual "derived - IsThree 1" 1 (length results)
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.IsThree
      [s|
        Y where Y = glean.test.IsThree X; X = 4
      |]
    print results
    assertEqual "derived - IsThree 2" 0 (length results)
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.IsThree
      [s|
        glean.test.IsThree 3
      |]
    print results
    assertEqual "derived - IsThree 3" 1 (length results)
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.IsThree
      [s|
        glean.test.IsThree 4
      |]
    print results
    assertEqual "derived - IsThree 4" 0 (length results)
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.IsThree
      [s|
        X = 3; glean.test.IsThree X
      |]
    print results
    assertEqual "derived - IsThree 5" 1 (length results)
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.IsThree
      [s|
        X = 4; glean.test.IsThree X
      |]
    print results
    assertEqual "derived - IsThree 6" 0 (length results)
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.IsGlean
      [s|
        glean.test.IsGlean "gl"..
      |]
    print results
    assertEqual "derived - IsGlean 1" 1 (length results)
  
    -- Tests with RevStringPair
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.RevStringPair
      [s|
        glean.test.RevStringPair _
      |]
    print results
    assertBool "derived 1" $ length results == 6
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.RevStringPair
      [s|
        glean.test.RevStringPair {"a",_}
      |]
    print results
    assertBool "derived 2" $ length results == 1
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.RevStringPair
      [s|
        glean.test.RevStringPair {_,"a"}
      |]
    print results
    assertBool "derived 3" $ case results of
      [one] | Just key <- getFactKey one ->
        "a" == Glean.Test.revStringPair_key_snd key
      _otherwise -> False
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.RevStringPair
      [s|
        glean.test.RevStringPair {X,X}
      |]
    print results
    assertBool "derived 4" $ length results == 2
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.RevRevStringPair
      [s|
        glean.test.RevRevStringPair _
      |]
    print results
    assertBool "derived 5" $ length results == 6
  
    -- The Ids of derived fact better be distinct, even when we page the results
    assertBool "derived 5: distinct Ids" $ length (nub (map getId results)) == 6
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.RevRevStringPair
      [s|
        glean.test.RevRevStringPair {"a",_}
      |]
    print results
    assertBool "derived 6" $ case results of
      [one] | Just key <- getFactKey one ->
        "a" == Glean.Test.revRevStringPair_key_fst key
      _otherwise -> False
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.ReflStringPair
      [s|
        glean.test.ReflStringPair "x"..
      |]
    print results
    assertBool "derived 7" $ length results == 1
  
    -- check that we can recursively expand the results of a derived query
    results <- runQuery_ env repo $ modify $ recursive $
      angle @Glean.Test.DualStringPair
      [s|
        glean.test.DualStringPair _
      |]
    print results
    assertBool "derived 8" $ (==4) $ length
      [ fst
      | Just key <- map getFactKey results
      , Just fst <- [getFactKey $ Glean.Test.dualStringPair_key_fst key] ]
  
    -- a pattern on the LHS of a statement that arises due to
    -- substitution of a derived predicate
    results <- runQuery_ env repo $ modify $ recursive $
      angle @Glean.Test.DualStringPair
      [s|
        P = glean.test.StringPair _; glean.test.DualStringPair {P,_}
      |]
    print results
    assertBool "derived 10" $ (==4) $ length results
  
    -- More complicated matching
    results <- runQuery_ env repo $ modify $ recursive $
      angle @Glean.Test.RevStringPair
      [s|
        glean.test.StringPair X;
        glean.test.RevStringPair X
        # generates all RevStringPairs and matches each one against X
      |]
    print results
    assertEqual "derived match filter" 4 (length results)
  
    -- deriving a key-only predicate from a key-value predicate
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.DerivedKeyValue
      [s|
        glean.test.DerivedKeyValue _
      |]
    print results
    assertEqual "angle - derived key/value 1" 3 (length results)
  
    -- match on the key
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.DerivedKeyValue
      [s|
        glean.test.DerivedKeyValue { kstring = "foo" }
      |]
    print results
    assertEqual "angle - derived key/value 2" 2 (length results)
  
    -- match on the value
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.DerivedKeyValue
      [s|
        glean.test.DerivedKeyValue { vnat = 5 }
      |]
    print results
    assertEqual "angle - derived key/value 3" 1 (length results)
  
    -- deriving a key-value predicate from a key-only predicate which is
    -- derived from a key-value predicate
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.DerivedKeyValue2
      [s|
        glean.test.DerivedKeyValue2 _
      |]
    print results
    assertEqual "angle - derived key/value 4" 3 (length results)
  
    -- match on the key
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.DerivedKeyValue2
      [s|
        glean.test.DerivedKeyValue2 { kstring = "foo" } -> _
      |]
    print results
    assertEqual "angle - derived key/value 5" 2 (length results)
  
    -- match on the value
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.DerivedKeyValue2
      [s|
        glean.test.DerivedKeyValue2 _ -> { vnat = 5 }
      |]
    print results
    assertEqual "angle - derived key/value 6" 1 (length results)
  
    -- a derived predicate defined as a record, not a tuple
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.RevStringPairRec
      [s|
        glean.test.RevStringPairRec {_, "a"}
      |]
    print results
    assertBool "angle - derived with record" $ length results == 1
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.RevStringPairRec
      [s|
        glean.test.RevStringPairRec { snd = "a" }
      |]
    print results
    assertBool "angle - derived with record 2" $ length results == 1
  
    -- tests that require inserting derived facts into the temporary environment
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.RevStringPair
      [s|
        X where X = glean.test.RevStringPair _
      |]
    print results
    assertEqual "angle - tmp environment" 6 (length results)
  
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.RevStringPair
      [s|
         P where
         P = glean.test.RevStringPair {A,"a"};
         glean.test.RevStringPair {_,A}
      |]
    print results
    assertEqual "angle - tmp environment 2" 1 (length results)
  
    results <- runQuery_ env repo $ modify $ recursive $
      angle @Glean.Test.RevStringPairs
      [s|
        glean.test.RevStringPairs { x = "a" }
      |]
    print results
    assertEqual "angle - tmp environment 3" 1 (length results)
  
    -- non-bindable or-pattern in the key of a derived predicate. This
    -- is a tricky case to handle: the query engine will have to expand
    -- and substitute the derived predicate twice.
    results <- runQuery_ env repo $ modify $ recursive $
      angle @Glean.Test.RevStringPairs
      [s|
        glean.test.RevStringPairs { x = "a".. | "b".. }
      |]
    print results
    assertEqual "angle - derived + or pattern" 2 (length results)
  
    results <- runQuery_ env repo $ modify $ recursive $
      angleData @[Text]
      [s|
        [X1, X2] where glean.test.RevStringPair{X1, X2};
      |]
    print results
    assertEqual "angle - derived returning array" 6 (length results)
  
    -- sum type in the derived predicate head
    results <- runQuery_ env repo $ modify $ recursive $
      angle @Glean.Test.MatchOneAlt
      [s|
        glean.test.MatchOneAlt { { tue = 3 }, _ }
      |]
    print results
    assertEqual "angle - MatchOneAlt 1" 1 (length results)
  
    results <- runQuery_ env repo $ modify $ recursive $
      angle @Glean.Test.MatchOneAlt
      [s|
        glean.test.MatchOneAlt { { mon = 3 }, _ }
      |]
    print results
    assertEqual "angle - MatchOneAlt 2" 0 (length results)
  
    -- Test for a bug
    r <- runQuery_ env repo $ modify $ angle @Glean.Test.LeftOr $
      [s|
        glean.test.LeftOr { x = "c".. }
      |]
    print r
    assertEqual "angle - LeftOr" 1 (length r)
  
    r <- runQuery_ env repo $ modify $ angle @Glean.Test.LeftOr2 $
      [s|
        glean.test.LeftOr2 { x = "c".. }
      |]
    print r
    assertEqual "angle - LeftOr2" 1 (length r)
  
    -- Test for another bug
    results <- runQuery_ env repo $ modify $ recursive $
      angle @Glean.Test.RevStringPairs
      [s|
        glean.test.RevStringPairs { X, {_, X} }
      |]
    print results
    assertEqual "angle - RevStringPairs " 6 (length results)
  
    -- flipping statements
    r <- runQuery_ env repo $ modify $ angle @Glean.Test.Unbound2
      [s|
        glean.test.Unbound2 { _, "e" }
      |]
    print r
    assertEqual "angle - flip statement 1" 1 (length r)
  
    r <- runQuery_ env repo $ modify $ angle @Glean.Test.Unbound2
      [s|
        glean.test.Unbound2 { "e", _ }
      |]
    print r
    assertEqual "angle - flip statement 2" 1 (length r)
  
    -- test for repeated variables in the derived predicate
    r <- runQuery_ env repo $ modify $ angle @Glean.Test.SameString
      [s|
        glean.test.SameString { "a", "b" }
      |]
    print r
    assertEqual "angle - repeated vars 1" 0 (length r)
  
    r <- runQuery_ env repo $ modify $ angle @Glean.Test.SameString
      [s|
        glean.test.SameString { "a", "a" }
      |]
    print r
    assertEqual "angle - repeated vars 2" 1 (length r)
  
    -- unification
    r <- runQuery_ env repo $ modify $ angle @Glean.Test.Unbound2
      [s|
        glean.test.Unbound2 _
      |]
    print r
    assertEqual "angle - unification" 2 (length r)
  
    -- or-statements and derived predicates
    r <- runQuery_ env repo $ modify $ angleData @Text
      [s|
        N where
          glean.test.RevStringPair { N, _ } |
          glean.test.RevRevStringPair { N, _ }
      |]
    assertEqual "angle - or-statements derived" 8 (length (nub r))
  