-- Copyright (c) Facebook, Inc. and its affiliates.
--
-- This source code is licensed under the MIT license found in the
-- LICENSE file in the root directory of this source tree.
--
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
module Retrie.GroundTerms
  ( GroundTerms
  , groundTerms
  ) where

import Data.HashSet (HashSet)
import qualified Data.HashSet as HashSet

import Retrie.ExactPrint
import Retrie.GHC
import Retrie.Quantifiers
import Retrie.SYB
import Retrie.Types

------------------------------------------------------------------------

-- | 'Ground Terms' are variables in the pattern that are not quantifiers.
-- We use a set of ground terms to save time during matching by filtering out
-- files which do not contain all the terms. We store one set of terms per
-- pattern because when filtering we must take care to only filter files which
-- do not match any of the patterns.
--
-- Example:
--
-- Patterns of 'forall x. foo (bar x) = ...' and 'forall y. baz (quux y) = ...'
--
-- groundTerms = [{'foo', 'bar'}, {'baz', 'quux'}]
--
-- Files will be found by unioning results of these commands:
--
-- grep -R --include "*.hs" -l foo dir | xargs grep -l bar
-- grep -R --include "*.hs" -l baz dir | xargs grep -l quux
--
-- If there are no ground terms (e.g. 'forall f x y. f x y = f y x')
-- we fall back to 'find dir -iname "*.hs"'. This case seems pathological.
type GroundTerms = HashSet String

groundTerms :: Data k => Query k v -> GroundTerms
groundTerms :: forall k v. Data k => Query k v -> GroundTerms
groundTerms Query{v
Quantifiers
Annotated k
qQuantifiers :: Quantifiers
qPattern :: Annotated k
qResult :: v
qQuantifiers :: forall ast v. Query ast v -> Quantifiers
qPattern :: forall ast v. Query ast v -> Annotated ast
qResult :: forall ast v. Query ast v -> v
..} = [String] -> GroundTerms
forall a. (Eq a, Hashable a) => [a] -> HashSet a
HashSet.fromList ([String] -> GroundTerms) -> [String] -> GroundTerms
forall a b. (a -> b) -> a -> b
$ k -> [String]
GenericQ [String]
go (k -> [String]) -> k -> [String]
forall a b. (a -> b) -> a -> b
$ Annotated k -> k
forall ast. Annotated ast -> ast
astA Annotated k
qPattern
  where
    go :: GenericQ [String]
    go :: GenericQ [String]
go a
x
      -- 'x' contains a quantifier, so split it into subtrees
      | (Bool -> Bool -> Bool) -> GenericQ Bool -> GenericQ Bool
forall r. (r -> r -> r) -> GenericQ r -> GenericQ r
everything Bool -> Bool -> Bool
(||) a -> Bool
GenericQ Bool
isQuantifier a
x = [[String]] -> [String]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[String]] -> [String]) -> [[String]] -> [String]
forall a b. (a -> b) -> a -> b
$ GenericQ [String] -> a -> [[String]]
forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
forall u. (forall d. Data d => d -> u) -> a -> [u]
gmapQ d -> [String]
GenericQ [String]
go a
x
      -- 'x' doesn't contain a quantifier, and can be exactprinted, so return
      -- the result of exactprinting
      | strs :: [String]
strs@(String
_:[String]
_) <- a -> [String]
GenericQ [String]
printer a
x = [String]
strs
      -- 'x' cannot be exactprinted, so recurse to find a printable child
      | Bool
otherwise = [[String]] -> [String]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[String]] -> [String]) -> [[String]] -> [String]
forall a b. (a -> b) -> a -> b
$ GenericQ [String] -> a -> [[String]]
forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
forall u. (forall d. Data d => d -> u) -> a -> [u]
gmapQ d -> [String]
GenericQ [String]
go a
x

    isQuantifier :: GenericQ Bool
    isQuantifier :: GenericQ Bool
isQuantifier = Bool -> (HsExpr GhcPs -> Bool) -> a -> Bool
forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
mkQ Bool
False HsExpr GhcPs -> Bool
exprIsQ (a -> Bool) -> (HsType GhcPs -> Bool) -> a -> Bool
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` HsType GhcPs -> Bool
tyIsQ

    -- returns 'True' if expression is a var that is a quantifier
    exprIsQ :: HsExpr GhcPs -> Bool
    exprIsQ :: HsExpr GhcPs -> Bool
exprIsQ HsExpr GhcPs
e | Just (L SrcSpanAnnN
_ RdrName
v) <- HsExpr GhcPs -> Maybe (LIdP GhcPs)
forall p. HsExpr p -> Maybe (LIdP p)
varRdrName HsExpr GhcPs
e = RdrName -> Quantifiers -> Bool
isQ RdrName
v Quantifiers
qQuantifiers
    exprIsQ HsExpr GhcPs
_ = Bool
False

    -- returns 'True' if type is a tyvar that is a quantifier
    tyIsQ :: HsType GhcPs -> Bool
    tyIsQ :: HsType GhcPs -> Bool
tyIsQ HsType GhcPs
ty | Just (L SrcSpanAnnN
_ RdrName
v) <- HsType GhcPs -> Maybe (LIdP GhcPs)
forall p. HsType p -> Maybe (LIdP p)
tyvarRdrName HsType GhcPs
ty = RdrName -> Quantifiers -> Bool
isQ RdrName
v Quantifiers
qQuantifiers
    tyIsQ HsType GhcPs
_ = Bool
False

    -- exactprinter that only works for expressions and types
    printer :: GenericQ [String]
    printer :: GenericQ [String]
printer = [String]
-> (LocatedAn AnnListItem (HsExpr GhcPs) -> [String])
-> a
-> [String]
forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
mkQ [] LHsExpr GhcPs -> [String]
LocatedAn AnnListItem (HsExpr GhcPs) -> [String]
printExpr (a -> [String])
-> (LocatedAn AnnListItem (HsType GhcPs) -> [String])
-> a
-> [String]
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` LHsType GhcPs -> [String]
LocatedAn AnnListItem (HsType GhcPs) -> [String]
printTy

    printExpr :: LHsExpr GhcPs -> [String]
    printExpr :: LHsExpr GhcPs -> [String]
printExpr LHsExpr GhcPs
e = [LocatedAn AnnListItem (HsExpr GhcPs) -> String
forall ast. ExactPrint ast => ast -> String
exactPrint (LocatedAn AnnListItem (HsExpr GhcPs)
-> DeltaPos -> LocatedAn AnnListItem (HsExpr GhcPs)
forall t a. Default t => LocatedAn t a -> DeltaPos -> LocatedAn t a
setEntryDP LHsExpr GhcPs
LocatedAn AnnListItem (HsExpr GhcPs)
e (Int -> DeltaPos
SameLine Int
0))]

    printTy :: LHsType GhcPs -> [String]
    printTy :: LHsType GhcPs -> [String]
printTy LHsType GhcPs
t = [LocatedAn AnnListItem (HsType GhcPs) -> String
forall ast. ExactPrint ast => ast -> String
exactPrint (LocatedAn AnnListItem (HsType GhcPs)
-> DeltaPos -> LocatedAn AnnListItem (HsType GhcPs)
forall t a. Default t => LocatedAn t a -> DeltaPos -> LocatedAn t a
setEntryDP LHsType GhcPs
LocatedAn AnnListItem (HsType GhcPs)
t (Int -> DeltaPos
SameLine Int
0))]