From 598e81bb845db4dd8fa6f1d51bf7811b24248191 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Fri, 8 Mar 2024 09:51:11 -0800
Subject: [PATCH] Replace internal reverse_postorder by a stable one

---
 .../codegen_cprover_gotoc/codegen/block.rs    | 25 ++++++++++++++++++-
 .../codegen_cprover_gotoc/codegen/function.rs |  6 ++---
 .../context/current_fn.rs                     |  9 -------
 3 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/kani-compiler/src/codegen_cprover_gotoc/codegen/block.rs b/kani-compiler/src/codegen_cprover_gotoc/codegen/block.rs
index 692fdb9c385b..bbf5739af3e3 100644
--- a/kani-compiler/src/codegen_cprover_gotoc/codegen/block.rs
+++ b/kani-compiler/src/codegen_cprover_gotoc/codegen/block.rs
@@ -2,7 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
 use crate::codegen_cprover_gotoc::GotocCtx;
-use stable_mir::mir::{BasicBlock, BasicBlockIdx};
+use stable_mir::mir::{BasicBlock, BasicBlockIdx, Body};
+use std::collections::HashSet;
 use tracing::debug;
 
 pub fn bb_label(bb: BasicBlockIdx) -> String {
@@ -72,3 +73,25 @@ impl<'tcx> GotocCtx<'tcx> {
         }
     }
 }
+
+pub fn reverse_postorder(body: &Body) -> impl Iterator<Item = BasicBlockIdx> {
+    postorder(body, 0, &mut HashSet::with_capacity(body.blocks.len())).into_iter().rev()
+}
+
+fn postorder(
+    body: &Body,
+    bb: BasicBlockIdx,
+    visited: &mut HashSet<BasicBlockIdx>,
+) -> Vec<BasicBlockIdx> {
+    if visited.contains(&bb) {
+        return vec![];
+    }
+    visited.insert(bb);
+
+    let mut result = vec![];
+    for succ in body.blocks[bb].terminator.successors() {
+        result.append(&mut postorder(body, succ, visited));
+    }
+    result.push(bb);
+    result
+}
diff --git a/kani-compiler/src/codegen_cprover_gotoc/codegen/function.rs b/kani-compiler/src/codegen_cprover_gotoc/codegen/function.rs
index 78c7e2b6cbd6..7cf4b979f407 100644
--- a/kani-compiler/src/codegen_cprover_gotoc/codegen/function.rs
+++ b/kani-compiler/src/codegen_cprover_gotoc/codegen/function.rs
@@ -3,10 +3,10 @@
 
 //! This file contains functions related to codegenning MIR functions into gotoc
 
+use crate::codegen_cprover_gotoc::codegen::block::reverse_postorder;
 use crate::codegen_cprover_gotoc::GotocCtx;
 use cbmc::goto_program::{Expr, Stmt, Symbol};
 use cbmc::InternString;
-use rustc_middle::mir::traversal::reverse_postorder;
 use stable_mir::mir::mono::Instance;
 use stable_mir::mir::{Body, Local};
 use stable_mir::ty::{RigidTy, TyKind};
@@ -67,9 +67,7 @@ impl<'tcx> GotocCtx<'tcx> {
             self.codegen_declare_variables(&body);
 
             // Get the order from internal body for now.
-            let internal_body = self.current_fn().body_internal();
-            reverse_postorder(internal_body)
-                .for_each(|(bb, _)| self.codegen_block(bb.index(), &body.blocks[bb.index()]));
+            reverse_postorder(&body).for_each(|bb| self.codegen_block(bb, &body.blocks[bb]));
 
             let loc = self.codegen_span_stable(instance.def.span());
             let stmts = self.current_fn_mut().extract_block();
diff --git a/kani-compiler/src/codegen_cprover_gotoc/context/current_fn.rs b/kani-compiler/src/codegen_cprover_gotoc/context/current_fn.rs
index adac294e30eb..5a542ca07cd2 100644
--- a/kani-compiler/src/codegen_cprover_gotoc/context/current_fn.rs
+++ b/kani-compiler/src/codegen_cprover_gotoc/context/current_fn.rs
@@ -4,7 +4,6 @@
 use crate::codegen_cprover_gotoc::GotocCtx;
 use cbmc::goto_program::Stmt;
 use cbmc::InternedString;
-use rustc_middle::mir::Body as BodyInternal;
 use rustc_middle::ty::Instance as InstanceInternal;
 use rustc_smir::rustc_internal;
 use stable_mir::mir::mono::Instance;
@@ -21,8 +20,6 @@ pub struct CurrentFnCtx<'tcx> {
     instance: Instance,
     /// The crate this function is from
     krate: String,
-    /// The MIR for the current instance. This is using the internal representation.
-    mir: &'tcx BodyInternal<'tcx>,
     /// The current instance. This is using the internal representation.
     instance_internal: InstanceInternal<'tcx>,
     /// A list of local declarations used to retrieve MIR component types.
@@ -53,7 +50,6 @@ impl<'tcx> CurrentFnCtx<'tcx> {
         Self {
             block: vec![],
             instance,
-            mir: gcx.tcx.instance_mir(instance_internal.def),
             instance_internal,
             krate: instance.def.krate().name,
             locals,
@@ -94,11 +90,6 @@ impl<'tcx> CurrentFnCtx<'tcx> {
         self.instance
     }
 
-    /// The internal MIR for the function we are currently compiling using internal APIs.
-    pub fn body_internal(&self) -> &'tcx BodyInternal<'tcx> {
-        self.mir
-    }
-
     /// The name of the function we are currently compiling
     pub fn name(&self) -> String {
         self.name.clone()