From b33129409a691a1ab581ca7fc65681a47c953cd6 Mon Sep 17 00:00:00 2001
From: Khyber Sen <kkysen@gmail.com>
Date: Mon, 16 Dec 2024 02:08:56 -0800
Subject: [PATCH] rewriter: add `--enable-dav1d_get_picture-post-condition`
 option, on by default

We still want to be able to run the rewriter on `dav1d`
normally without post conditions while this is still a WIP.
Or on Arm, where IA2 still lacks support for post conditions.
---
 tools/rewriter/GenCallAsm.cpp     | 3 ++-
 tools/rewriter/GenCallAsm.h       | 2 ++
 tools/rewriter/SourceRewriter.cpp | 8 ++++++++
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/tools/rewriter/GenCallAsm.cpp b/tools/rewriter/GenCallAsm.cpp
index e7db75eac..782a8887f 100644
--- a/tools/rewriter/GenCallAsm.cpp
+++ b/tools/rewriter/GenCallAsm.cpp
@@ -1053,9 +1053,10 @@ std::string emit_asm_wrapper(AbiSignature &sig,
   // The name is `${target_name}_post_condition`,
   // and we only do this for `dav1d_get_picture`.
   std::optional<std::string> target_post_condition_name = std::nullopt;
-  if (target_name && *target_name == "dav1d_get_picture") {
+  if (enable_dav1d_get_picture_post_condition && target_name && *target_name == "dav1d_get_picture") {
     target_post_condition_name = *target_name + "_post_condition";
   }
+
   if (target_post_condition_name) {
     emit_post_condition_fn_call(aw, arch, *target_post_condition_name);
   }
diff --git a/tools/rewriter/GenCallAsm.h b/tools/rewriter/GenCallAsm.h
index 88c004096..91f5a48b5 100644
--- a/tools/rewriter/GenCallAsm.h
+++ b/tools/rewriter/GenCallAsm.h
@@ -20,6 +20,8 @@ enum class Arch {
     Aarch64, X86
 };
 
+extern bool enable_dav1d_get_picture_post_condition;
+
 // Generates a wrapper for a function named \p name with the signature \p sig.
 // The WrapperKind parameter \p kind determines the type of call which may
 // affect the order of operations and the layout of the wrapper stack frame.
diff --git a/tools/rewriter/SourceRewriter.cpp b/tools/rewriter/SourceRewriter.cpp
index 3d183eac3..346f34d29 100644
--- a/tools/rewriter/SourceRewriter.cpp
+++ b/tools/rewriter/SourceRewriter.cpp
@@ -98,10 +98,17 @@ static llvm::cl::opt<std::string>
                        llvm::cl::cat(SourceRewriterCategory),
                        llvm::cl::desc("<prefix for output files>"));
 
+static llvm::cl::opt<bool>
+    EnableDav1dGetPicturePostCondition("enable-dav1d_get_picture-post-condition",
+                                       llvm::cl::init(true),
+                                       llvm::cl::cat(SourceRewriterCategory),
+                                       llvm::cl::desc("enable calling the post condition function hardcoded for dav1d_get_picture"));
+
 static Arch Target;
 static std::string RootDirectory;
 static std::string OutputDirectory;
 static std::string OutputPrefix;
+bool enable_dav1d_get_picture_post_condition = true;
 
 // Map each translation unit's filename to its pkey.
 static std::map<Filename, Pkey> file_pkeys;
@@ -1098,6 +1105,7 @@ int main(int argc, const char **argv) {
   RootDirectory = RootDirectoryOption;
   OutputDirectory = OutputDirectoryOption;
   OutputPrefix = OutputPrefixOption;
+  enable_dav1d_get_picture_post_condition = EnableDav1dGetPicturePostCondition;
 
   RefactoringTool tool(options_parser.getCompilations(),
                        options_parser.getSourcePathList());