-- 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
qResult :: forall ast v. Query ast v -> v
qPattern :: forall ast v. Query ast v -> Annotated ast
qQuantifiers :: forall ast v. Query ast v -> Quantifiers
qResult :: v
qPattern :: Annotated k
qQuantifiers :: Quantifiers
..} = forall a. (Eq a, Hashable a) => [a] -> HashSet a
HashSet.fromList forall a b. (a -> b) -> a -> b
$ GenericQ [String]
go forall a b. (a -> b) -> a -> b
$ 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
      | forall r. (r -> r -> r) -> GenericQ r -> GenericQ r
everything Bool -> Bool -> Bool
(||) GenericQ Bool
isQuantifier a
x = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$ forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
gmapQ 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]
_) <- GenericQ [String]
printer a
x = [String]
strs
      -- 'x' cannot be exactprinted, so recurse to find a printable child
      | Bool
otherwise = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$ forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
gmapQ GenericQ [String]
go a
x

    isQuantifier :: GenericQ Bool
    isQuantifier :: GenericQ Bool
isQuantifier = forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
mkQ Bool
False HsExpr GhcPs -> Bool
exprIsQ forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`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) <- 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) <- 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 = forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
mkQ [] LHsExpr GhcPs -> [String]
printExpr forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`extQ` LHsType GhcPs -> [String]
printTy

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

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