From 8ab8b7d6c704226eed3af278969da2c3751bf84e Mon Sep 17 00:00:00 2001
From: Fabian Ruffy <5960321+fruffy@users.noreply.github.com>
Date: Thu, 31 Oct 2024 08:09:00 -0400
Subject: [PATCH] Run typechecking after front and mid end. (#4834)

Signed-off-by: fruffy <fruffy@nyu.edu>
---
 .../p4tools/common/compiler/compiler_target.cpp     | 13 ++++++++-----
 backends/p4tools/common/compiler/midend.cpp         | 12 ++++++++++--
 backends/p4tools/common/compiler/midend.h           |  6 +++++-
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/backends/p4tools/common/compiler/compiler_target.cpp b/backends/p4tools/common/compiler/compiler_target.cpp
index 7926f309a58..96a05c29fb7 100644
--- a/backends/p4tools/common/compiler/compiler_target.cpp
+++ b/backends/p4tools/common/compiler/compiler_target.cpp
@@ -9,7 +9,6 @@
 #include "frontends/common/parseInput.h"
 #include "frontends/common/parser_options.h"
 #include "frontends/p4/frontend.h"
-#include "lib/compile_context.h"
 #include "lib/error.h"
 
 namespace P4::P4Tools {
@@ -27,7 +26,7 @@ CompilerResultOrError CompilerTarget::runCompiler(const CompilerOptions &options
 CompilerResultOrError CompilerTarget::runCompiler(const CompilerOptions &options,
                                                   std::string_view toolName,
                                                   const std::string &source) {
-    const auto *program = P4::parseP4String(source, options.langVersion);
+    const auto *program = parseP4String(source, options.langVersion);
     if (program == nullptr) {
         return std::nullopt;
     }
@@ -57,7 +56,7 @@ CompilerResultOrError CompilerTarget::runCompilerImpl(const CompilerOptions &opt
 }
 
 const IR::P4Program *CompilerTarget::runParser(const ParserOptions &options) {
-    const auto *program = P4::parseP4File(options);
+    const auto *program = parseP4File(options);
     if (errorCount() > 0) {
         return nullptr;
     }
@@ -66,12 +65,16 @@ const IR::P4Program *CompilerTarget::runParser(const ParserOptions &options) {
 
 const IR::P4Program *CompilerTarget::runFrontend(const CompilerOptions &options,
                                                  const IR::P4Program *program) const {
-    P4::P4COptionPragmaParser optionsPragmaParser;
-    program->apply(P4::ApplyOptionsPragmas(optionsPragmaParser));
+    P4COptionPragmaParser optionsPragmaParser;
+    program->apply(ApplyOptionsPragmas(optionsPragmaParser));
 
     auto frontEnd = mkFrontEnd();
     frontEnd.addDebugHook(options.getDebugHook());
     program = frontEnd.run(options, program);
+    ReferenceMap refMap;
+    TypeMap typeMap;
+    // Perform a last round of type checking.
+    program = program->apply(TypeChecking(&refMap, &typeMap, true));
     if ((program == nullptr) || errorCount() > 0) {
         return nullptr;
     }
diff --git a/backends/p4tools/common/compiler/midend.cpp b/backends/p4tools/common/compiler/midend.cpp
index f356d62119a..6d58f28ca76 100644
--- a/backends/p4tools/common/compiler/midend.cpp
+++ b/backends/p4tools/common/compiler/midend.cpp
@@ -156,9 +156,18 @@ void MidEnd::addDefaultPasses() {
         new P4::EliminateTuples(&typeMap),
         new P4::ConstantFolding(&typeMap),
         new P4::SimplifyControlFlow(&typeMap),
+        // Perform a last round of type-checking before passes which do not type-check begin.
+        new P4::TypeChecking(&refMap, &typeMap, true),
+    });
+    addNonTypeCheckingPasses();
+}
+
+void MidEnd::addNonTypeCheckingPasses() {
+    addPasses({
         // Simplify header stack assignments with runtime indices into conditional statements.
         new P4::HSIndexSimplifier(&typeMap),
-        // Convert Type_Varbits into a type that contains information about the assigned width.
+        // Convert Type_Varbits into a type that contains information about the assigned
+        // width.
         new ConvertVarbits(),
         // Convert any StructExpressions with Type_Header into a HeaderExpression.
         new ConvertStructExpr(&typeMap),
@@ -166,5 +175,4 @@ void MidEnd::addDefaultPasses() {
         new P4::CastBooleanTableKeys(),
     });
 }
-
 }  // namespace P4::P4Tools
diff --git a/backends/p4tools/common/compiler/midend.h b/backends/p4tools/common/compiler/midend.h
index 95b3683a2d3..f5c6328db9d 100644
--- a/backends/p4tools/common/compiler/midend.h
+++ b/backends/p4tools/common/compiler/midend.h
@@ -64,7 +64,11 @@ class MidEnd : public PassManager {
 
     /// Add the list of default passes to the mid end. This is not part of the initializer because
     /// some targets may add their own passes to the beginning of the pass list.
-    void addDefaultPasses();
+    virtual void addDefaultPasses();
+
+    /// Add passes that break type checking. These passes may involve  IR modifications the front
+    /// end type checker does not recognize.
+    virtual void addNonTypeCheckingPasses();
 };
 
 }  // namespace P4::P4Tools