Optimize required_directive, and add a directive_finder trait

Reviewed By: kassens

Differential Revision: D32378752

fbshipit-source-id: 0eaa5ea30b816fd03cdc638dd0f3654837957712

compiler/crates/graphql-ir/src/early_return_visitor.rs

@@ -0,0 +1,46 @@
+/*
+ * 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.
+ */
+
+pub trait EarlyReturnVisitor {
+ fn visit(selections: &[Selection]) {
+ let mut to_visit = Vec::from(selections);
+ while let Some(selection) = to_visit.pop() {
+ match selection {
+ Selection::FragmentSpread(s) => {}
+ Selection::InlineFragment(_) => todo!(),
+ Selection::LinkedField(_) => todo!(),
+ Selection::ScalarField(_) => todo!(),
+ Selection::Condition(_) => todo!(),
+ }
+ }
+ }
+
+ fn visit_fragment(
+ &mut self,
+ to_visit: &mut Vec<&Selection>,
+ fragment: &FragmentDefinition,
+ ) -> bool {
+ to_visit.push(fragment.selections);
+ false
+ }
+
+ fn visit_inline_fragment(
+ &mut self,
+ to_visit: &mut Vec<&Selection>,
+ fragment: &InlineFragment,
+ ) -> bool {
+ to_visit.push(fragment.selections);
+ }
+
+ fn visit_linkedField(
+ &mut self,
+ to_visit: &mut Vec<&Selection>,
+ fragment: &InlineFragment,
+ ) -> bool {
+ to_visit.push(fragment.selections);
+ }
+}

compiler/crates/relay-transforms/src/directive_finder.rs

@@ -0,0 +1,76 @@
+/*
+ * 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.
+ */
+
+use graphql_ir::{Directive, FragmentSpread, InlineFragment, LinkedField, Selection};
+
+/// A trait for implementating finding if a directive exists in an IR.
+/// Returns `true` if the directive is found.
+/// We can make it more general, if there are more similar use cases
+pub trait DirectiveFinder {
+ fn find(&mut self, mut to_visit: Vec<&Selection>) -> bool {
+ while let Some(selection) = to_visit.pop() {
+ match selection {
+ Selection::FragmentSpread(f) => {
+ if self.visit_fragment_spread(f) {
+ return true;
+ }
+ }
+ Selection::InlineFragment(f) => {
+ if self.visit_inline_fragment(&mut to_visit, f) {
+ return true;
+ }
+ }
+ Selection::LinkedField(f) => {
+ if self.visit_linked_field(&mut to_visit, f) {
+ return true;
+ }
+ }
+ Selection::ScalarField(f) => {
+ if self.visit_directives(&f.directives) {
+ return true;
+ }
+ }
+ Selection::Condition(c) => {
+ to_visit.extend(c.selections.iter());
+ }
+ }
+ }
+ false
+ }
+
+ fn visit_directive(&self, directive: &Directive) -> bool;
+
+ fn visit_directives(&self, directives: &[Directive]) -> bool {
+ directives.iter().any(|d| self.visit_directive(d))
+ }
+
+ fn visit_fragment_spread(&mut self, _fragment_spread: &FragmentSpread) -> bool;
+
+ fn visit_inline_fragment<'a>(
+ &mut self,
+ to_visit: &mut Vec<&'a Selection>,
+ fragment: &'a InlineFragment,
+ ) -> bool {
+ if self.visit_directives(&fragment.directives) {
+ return true;
+ }
+ to_visit.extend(fragment.selections.iter());
+ false
+ }
+
+ fn visit_linked_field<'a>(
+ &mut self,
+ to_visit: &mut Vec<&'a Selection>,
+ field: &'a LinkedField,
+ ) -> bool {
+ if self.visit_directives(&field.directives) {
+ return true;
+ }
+ to_visit.extend(field.selections.iter());
+ false
+ }
+}

compiler/crates/relay-transforms/src/lib.rs

@@ -19,6 +19,7 @@ mod client_extensions;
 mod connections;
 mod declarative_connection;
 mod defer_stream;
+mod directive_finder;
 mod errors;
 mod flatten;
 mod generate_data_driven_dependency_metadata;
@@ -58,7 +59,6 @@ mod unwrap_custom_directive_selection;
 mod util;
 mod validate_operation_variables;
 mod validations;
-
 use fnv::{FnvHashMap, FnvHashSet};
 use interner::{Intern, StringKey};
 use lazy_static::lazy_static;
@@ -90,6 +90,7 @@ pub use declarative_connection::transform_declarative_connection;
 pub use defer_stream::{
  transform_defer_stream, DeferDirective, StreamDirective, DEFER_STREAM_CONSTANTS,
 };
+pub use directive_finder::DirectiveFinder;
 pub use flatten::flatten;
 pub use generate_data_driven_dependency_metadata::{
  generate_data_driven_dependency_metadata, DATA_DRIVEN_DEPENDENCY_METADATA_KEY,

compiler/crates/relay-transforms/src/required_directive/mod.rs

@@ -18,6 +18,8 @@ use lazy_static::lazy_static;
 use requireable_field::{RequireableField, RequiredMetadata};
 use std::{borrow::Cow, mem, sync::Arc};
 
+use crate::DirectiveFinder;
+
 lazy_static! {
  pub static ref REQUIRED_DIRECTIVE_NAME: StringKey = "required".intern();
  pub static ref ACTION_ARGUMENT: StringKey = "action".intern();
@@ -73,6 +75,7 @@ struct RequiredDirective<'s> {
  current_node_required_children: FnvHashMap<StringKey, RequiredField>,
  required_children_map: FnvHashMap<StringKey, FnvHashMap<StringKey, RequiredField>>,
  enabled: bool,
+ required_directive_visitor: RequiredDirectiveVisitor<'s>,
 }
 
 impl<'program> RequiredDirective<'program> {
@@ -87,6 +90,10 @@ impl<'program> RequiredDirective<'program> {
  current_node_required_children: Default::default(),
  required_children_map: Default::default(),
  enabled,
+ required_directive_visitor: RequiredDirectiveVisitor {
+ program,
+ visited_fragments: Default::default(),
+ },
  }
  }
 
@@ -297,6 +304,9 @@ impl<'s> Transformer for RequiredDirective<'s> {
  &mut self,
  fragment: &FragmentDefinition,
  ) -> Transformed<FragmentDefinition> {
+ if !self.required_directive_visitor.visit_fragment(fragment) {
+ return Transformed::Keep;
+ }
  self.reset_state();
  self.parent_inline_fragment_directive = fragment
  .directives
@@ -322,6 +332,12 @@ impl<'s> Transformer for RequiredDirective<'s> {
  &mut self,
  operation: &OperationDefinition,
  ) -> Transformed<OperationDefinition> {
+ if !self
+ .required_directive_visitor
+ .find(operation.selections.iter().collect())
+ {
+ return Transformed::Keep;
+ }
  self.reset_state();
  let selections = self.transform_selections(&operation.selections);
  let directives = maybe_add_children_can_bubble_metadata_directive(
@@ -506,3 +522,35 @@ impl From<StringKey> for RequiredAction {
  }
  }
 }
+
+struct RequiredDirectiveVisitor<'s> {
+ program: &'s Program,
+ visited_fragments: FnvHashMap<StringKey, bool>,
+}
+
+impl<'s> DirectiveFinder for RequiredDirectiveVisitor<'s> {
+ fn visit_directive(&self, directive: &Directive) -> bool {
+ directive.name.item == *REQUIRED_DIRECTIVE_NAME
+ }
+
+ fn visit_fragment_spread(&mut self, fragment_spread: &graphql_ir::FragmentSpread) -> bool {
+ let fragment = self
+ .program
+ .fragment(fragment_spread.fragment.item)
+ .unwrap();
+ self.visit_fragment(fragment)
+ }
+}
+
+impl<'s> RequiredDirectiveVisitor<'s> {
+ fn visit_fragment(&mut self, fragment: &FragmentDefinition) -> bool {
+ if let Some(val) = self.visited_fragments.get(&fragment.name.item) {
+ return *val;
+ }
+ // Avoid dead loop in self-referencing fragments
+ self.visited_fragments.insert(fragment.name.item, false);
+ let result = self.find(fragment.selections.iter().collect());
+ self.visited_fragments.insert(fragment.name.item, result);
+ result
+ }
+}