AST fuzzer: limit unary expression wrapping

Counts the size of expressions in the AST to avoid applying unary
wrapping to expressions that have already gotten large.

Change-Id: I0868d6f2bb3c6aaf99efdfb9574327d0af420456
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/84900
Reviewed-by: Antonio Maiorano <amaiorano@google.com>
Auto-Submit: Alastair Donaldson <afdx@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Alastair Donaldson <afdx@google.com>
diff --git a/src/tint/fuzzers/tint_ast_fuzzer/BUILD.gn b/src/tint/fuzzers/tint_ast_fuzzer/BUILD.gn
index b6d205a..0db6c1d 100644
--- a/src/tint/fuzzers/tint_ast_fuzzer/BUILD.gn
+++ b/src/tint/fuzzers/tint_ast_fuzzer/BUILD.gn
@@ -41,6 +41,8 @@
     sources = [
       "cli.cc",
       "cli.h",
+      "expression_size.cc",
+      "expression_size.h",
       "fuzzer.cc",
       "mutation.cc",
       "mutation.h",
diff --git a/src/tint/fuzzers/tint_ast_fuzzer/CMakeLists.txt b/src/tint/fuzzers/tint_ast_fuzzer/CMakeLists.txt
index 25dd3fc..1a45897 100644
--- a/src/tint/fuzzers/tint_ast_fuzzer/CMakeLists.txt
+++ b/src/tint/fuzzers/tint_ast_fuzzer/CMakeLists.txt
@@ -38,6 +38,7 @@
         ../mersenne_twister_engine.h
         ../random_generator.h
         ../random_generator_engine.h
+        expression_size.h
         mutation.h
         mutation_finder.h
         mutation_finders/change_binary_operators.h
@@ -57,6 +58,7 @@
         ../mersenne_twister_engine.cc
         ../random_generator.cc
         ../random_generator_engine.cc
+        expression_size.cc
         mutation.cc
         mutation_finder.cc
         mutation_finders/change_binary_operators.cc
@@ -100,9 +102,10 @@
 # Add tests.
 if (${TINT_BUILD_TESTS})
     set(TEST_SOURCES
+            expression_size_test.cc
             mutations/change_binary_operator_test.cc
             mutations/replace_identifier_test.cc
-	    mutations/wrap_unary_operator_test.cc)
+	        mutations/wrap_unary_operator_test.cc)
 
     add_executable(tint_ast_fuzzer_unittests ${TEST_SOURCES})
 
diff --git a/src/tint/fuzzers/tint_ast_fuzzer/expression_size.cc b/src/tint/fuzzers/tint_ast_fuzzer/expression_size.cc
new file mode 100644
index 0000000..fe4a5c4
--- /dev/null
+++ b/src/tint/fuzzers/tint_ast_fuzzer/expression_size.cc
@@ -0,0 +1,46 @@
+// Copyright 2022 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/tint/fuzzers/tint_ast_fuzzer/expression_size.h"
+
+#include "src/tint/ast/traverse_expressions.h"
+
+namespace tint::fuzzers::ast_fuzzer {
+
+ExpressionSize::ExpressionSize(const Program& program) {
+  // By construction, all the children of an AST node are encountered before the
+  // node itself when iterating through a program's AST nodes. Computing
+  // expression sizes exploits this property: the size of a compound expression
+  // is computed based on the already-computed sizes of its sub-expressions.
+  for (const auto* node : program.ASTNodes().Objects()) {
+    const auto* expr_ast_node = node->As<ast::Expression>();
+    if (expr_ast_node == nullptr) {
+      continue;
+    }
+    size_t expr_size = 0;
+    diag::List empty;
+    ast::TraverseExpressions(expr_ast_node, empty,
+                             [&](const ast::Expression* expression) {
+                               if (expression == expr_ast_node) {
+                                 expr_size++;
+                                 return ast::TraverseAction::Descend;
+                               }
+                               expr_size += expr_to_size_.at(expression);
+                               return ast::TraverseAction::Skip;
+                             });
+    expr_to_size_[expr_ast_node] = expr_size;
+  }
+}
+
+}  // namespace tint::fuzzers::ast_fuzzer
diff --git a/src/tint/fuzzers/tint_ast_fuzzer/expression_size.h b/src/tint/fuzzers/tint_ast_fuzzer/expression_size.h
new file mode 100644
index 0000000..5a6d3dd
--- /dev/null
+++ b/src/tint/fuzzers/tint_ast_fuzzer/expression_size.h
@@ -0,0 +1,47 @@
+// Copyright 2022 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef SRC_TINT_FUZZERS_TINT_AST_FUZZER_EXPRESSION_SIZE_H_
+#define SRC_TINT_FUZZERS_TINT_AST_FUZZER_EXPRESSION_SIZE_H_
+
+#include <unordered_map>
+
+#include "src/tint/ast/expression.h"
+#include "src/tint/program.h"
+
+namespace tint::fuzzers::ast_fuzzer {
+
+/// This class computes the size of the subtree rooted at each expression in a
+/// program, and allows these sizes to be subsequently queried.
+class ExpressionSize {
+ public:
+  /// Initializes expression size information for the given program.
+  /// @param program - the program for which expression sizes will be computed;
+  ///     must remain in scope as long as this instance exists.
+  explicit ExpressionSize(const Program& program);
+
+  /// Returns the size of the subtree rooted at the given expression.
+  /// @param expression - the expression whose size should be returned.
+  /// @return the size of the subtree rooted at `expression`.
+  size_t operator()(const ast::Expression* expression) const {
+    return expr_to_size_.at(expression);
+  }
+
+ private:
+  std::unordered_map<const ast::Expression*, size_t> expr_to_size_;
+};
+
+}  // namespace tint::fuzzers::ast_fuzzer
+
+#endif  // SRC_TINT_FUZZERS_TINT_AST_FUZZER_EXPRESSION_SIZE_H_
diff --git a/src/tint/fuzzers/tint_ast_fuzzer/expression_size_test.cc b/src/tint/fuzzers/tint_ast_fuzzer/expression_size_test.cc
new file mode 100644
index 0000000..bb6e5e7
--- /dev/null
+++ b/src/tint/fuzzers/tint_ast_fuzzer/expression_size_test.cc
@@ -0,0 +1,70 @@
+// Copyright 2022 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/tint/fuzzers/tint_ast_fuzzer/expression_size.h"
+
+#include <string>
+
+#include "gtest/gtest.h"
+
+#include "src/tint/ast/binary_expression.h"
+#include "src/tint/ast/expression.h"
+#include "src/tint/ast/int_literal_expression.h"
+#include "src/tint/program.h"
+#include "src/tint/reader/wgsl/parser.h"
+
+namespace tint {
+namespace fuzzers {
+namespace ast_fuzzer {
+namespace {
+
+TEST(ExpressionSizeTest, Basic) {
+  std::string content = R"(
+    fn main() {
+      let a = (0 + 0) * (0 + 0);
+    }
+  )";
+  Source::File file("test.wgsl", content);
+  auto program = reader::wgsl::Parse(&file);
+  ASSERT_TRUE(program.IsValid()) << program.Diagnostics().str();
+
+  ExpressionSize expression_size(program);
+  for (const auto* node : program.ASTNodes().Objects()) {
+    const auto* expr = node->As<ast::Expression>();
+    if (expr == nullptr) {
+      continue;
+    }
+    if (expr->Is<ast::IntLiteralExpression>()) {
+      ASSERT_EQ(1, expression_size(expr));
+    } else {
+      const auto* binary_expr = expr->As<ast::BinaryExpression>();
+      ASSERT_TRUE(binary_expr != nullptr);
+      switch (binary_expr->op) {
+        case ast::BinaryOp::kAdd:
+          ASSERT_EQ(3, expression_size(expr));
+          break;
+        case ast::BinaryOp::kMultiply:
+          ASSERT_EQ(7, expression_size(expr));
+          break;
+        default:
+          FAIL();
+      }
+    }
+  }
+}
+
+}  // namespace
+}  // namespace ast_fuzzer
+}  // namespace fuzzers
+}  // namespace tint
diff --git a/src/tint/fuzzers/tint_ast_fuzzer/mutation_finders/wrap_unary_operators.cc b/src/tint/fuzzers/tint_ast_fuzzer/mutation_finders/wrap_unary_operators.cc
index da026f1..e12f44a 100644
--- a/src/tint/fuzzers/tint_ast_fuzzer/mutation_finders/wrap_unary_operators.cc
+++ b/src/tint/fuzzers/tint_ast_fuzzer/mutation_finders/wrap_unary_operators.cc
@@ -17,6 +17,7 @@
 #include <memory>
 #include <vector>
 
+#include "src/tint/fuzzers/tint_ast_fuzzer/expression_size.h"
 #include "src/tint/fuzzers/tint_ast_fuzzer/mutations/wrap_unary_operator.h"
 #include "src/tint/fuzzers/tint_ast_fuzzer/util.h"
 #include "src/tint/sem/expression.h"
@@ -26,12 +27,18 @@
 namespace fuzzers {
 namespace ast_fuzzer {
 
+namespace {
+const size_t kMaxExpressionSize = 100;
+}  // namespace
+
 MutationList MutationFinderWrapUnaryOperators::FindMutations(
     const tint::Program& program,
     NodeIdMap* node_id_map,
     ProbabilityContext* probability_context) const {
   MutationList result;
 
+  ExpressionSize expression_size(program);
+
   // Iterate through all ast nodes and for each expression node, try to wrap
   // the inside a valid unary operator based on the type of the expression.
   for (const auto* node : program.ASTNodes().Objects()) {
@@ -42,6 +49,10 @@
       continue;
     }
 
+    if (expression_size(expr_ast_node) > kMaxExpressionSize) {
+      continue;
+    }
+
     const auto* expr_sem_node =
         tint::As<sem::Expression>(program.Sem().Get(expr_ast_node));