partial refactor of function_body

rr_frontend/translation/src/arg_folder.rs

@@ -2,9 +2,19 @@ use rr_rustc_interface::middle::ty::visit::*;
 use rr_rustc_interface::middle::ty::{self, GenericArg, GenericArgKind, ParamConst, Ty, TyCtxt, TypeFolder};
 use rr_rustc_interface::type_ir::fold::{TypeFoldable, TypeSuperFoldable};
 
+/// Instantiate a type with arguments.
+/// The type may contain open region variables `ReVar`.
+pub fn ty_instantiate<'tcx>(x: Ty<'tcx>, tcx: TyCtxt<'tcx>, args: &[GenericArg<'tcx>]) -> Ty<'tcx> {
+    let mut folder = ArgFolder {
+        tcx,
+        args,
+        binders_passed: 0,
+    };
+    x.fold_with(&mut folder)
+}
+
 /// A version of the `ArgFolder` in `rr_rustc_interface::middle::src::ty::generic_args` that skips over
 /// `ReVar` (instead of triggering a bug).
-
 struct ArgFolder<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     args: &'a [GenericArg<'tcx>],
@@ -208,12 +218,3 @@ impl<'a, 'tcx> ArgFolder<'a, 'tcx> {
         ty::fold::shift_region(self.tcx, region, self.binders_passed)
     }
 }
-
-pub fn ty_instantiate<'tcx>(x: Ty<'tcx>, tcx: TyCtxt<'tcx>, args: &[GenericArg<'tcx>]) -> Ty<'tcx> {
-    let mut folder = ArgFolder {
-        tcx,
-        args,
-        binders_passed: 0,
-    };
-    x.fold_with(&mut folder)
-}

rr_frontend/translation/src/body/mod.rs

+// © 2023, The RefinedRust Developers and Contributors
+//
+// This Source Code Form is subject to the terms of the BSD-3-clause License.
+// If a copy of the BSD-3-clause license was not distributed with this
+// file, You can obtain one at https://opensource.org/license/bsd-3-clause/.
+
+use std::collections::{btree_map, BTreeMap, HashMap, HashSet};
+
+use log::{info, trace, warn};
+use radium::coq;
+use rr_rustc_interface::hir::def_id::DefId;
+use rr_rustc_interface::middle::mir::interpret::{ConstValue, ErrorHandled, Scalar};
+use rr_rustc_interface::middle::mir::tcx::PlaceTy;
+use rr_rustc_interface::middle::mir::{
+    BasicBlock, BasicBlockData, BinOp, Body, BorrowKind, Constant, ConstantKind, Local, LocalKind, Location,
+    Mutability, NonDivergingIntrinsic, Operand, Place, ProjectionElem, Rvalue, StatementKind, Terminator,
+    TerminatorKind, UnOp, VarDebugInfoContents,
+};
+use rr_rustc_interface::middle::ty::fold::TypeFolder;
+use rr_rustc_interface::middle::ty::{ConstKind, Ty, TyKind};
+use rr_rustc_interface::middle::{mir, ty};
+use rr_rustc_interface::{abi, ast, middle};
+use typed_arena::Arena;
+
+use crate::arg_folder::*;
+use crate::base::*;
+use crate::checked_op_analysis::CheckedOpLocalAnalysis;
+use crate::environment::borrowck::facts;
+use crate::environment::polonius_info::PoloniusInfo;
+use crate::environment::procedure::Procedure;
+use crate::environment::{dump_borrowck_info, polonius_info, Environment};
+use crate::inclusion_tracker::*;
+use crate::spec_parsers::parse_utils::ParamLookup;
+use crate::spec_parsers::verbose_function_spec_parser::{
+    ClosureMetaInfo, FunctionRequirements, FunctionSpecParser, VerboseFunctionSpecParser,
+};
+use crate::traits::{registry, resolution};
+use crate::{base, consts, procedures, regions, search, traits, types};
+
+pub mod signature;
+pub mod translator;
+
+/// Get the syntypes of function arguments for a procedure call.
+pub fn get_arg_syntypes_for_procedure_call<'tcx, 'def>(
+    tcx: ty::TyCtxt<'tcx>,
+    ty_translator: &types::LocalTX<'def, 'tcx>,
+    callee_did: DefId,
+    ty_params: &[ty::GenericArg<'tcx>],
+) -> Result<Vec<radium::SynType>, TranslationError<'tcx>> {
+    let caller_did = ty_translator.get_proc_did();
+
+    // Get the type of the callee, fully instantiated
+    let full_ty: ty::EarlyBinder<Ty<'tcx>> = tcx.type_of(callee_did);
+    let full_ty = full_ty.instantiate(tcx, ty_params);
+
+    // We create a dummy scope in order to make the lifetime lookups succeed, since we only want
+    // syntactic types here.
+    // Since we do the substitution of the generics above, we should translate generics and
+    // traits in the caller's scope.
+    let scope = ty_translator.scope.borrow();
+    let param_env: ty::ParamEnv<'tcx> = tcx.param_env(scope.did);
+    let callee_state = types::CalleeState::new(&param_env, &scope.generic_scope);
+    let mut dummy_state = types::STInner::CalleeTranslation(callee_state);
+
+    let mut syntypes = Vec::new();
+    match full_ty.kind() {
+        ty::TyKind::FnDef(_, _) => {
+            let sig = full_ty.fn_sig(tcx);
+            for ty in sig.inputs().skip_binder() {
+                let st = ty_translator.translator.translate_type_to_syn_type(*ty, &mut dummy_state)?;
+                syntypes.push(st);
+            }
+        },
+        ty::TyKind::Closure(_, args) => {
+            let clos_args = args.as_closure();
+            let sig = clos_args.sig();
+            let pre_sig = sig.skip_binder();
+            // we also need to add the closure argument here
+
+            let tuple_ty = clos_args.tupled_upvars_ty();
+            match clos_args.kind() {
+                ty::ClosureKind::Fn | ty::ClosureKind::FnMut => {
+                    syntypes.push(radium::SynType::Ptr);
+                },
+                ty::ClosureKind::FnOnce => {
+                    let st =
+                        ty_translator.translator.translate_type_to_syn_type(tuple_ty, &mut dummy_state)?;
+                    syntypes.push(st);
+                },
+            }
+            for ty in pre_sig.inputs() {
+                let st = ty_translator.translator.translate_type_to_syn_type(*ty, &mut dummy_state)?;
+                syntypes.push(st);
+            }
+        },
+        _ => unimplemented!(),
+    }
+
+    Ok(syntypes)
+}
+
+// solve the constraints for the new_regions
+// we first identify what kinds of constraints these new regions are subject to
+#[derive(Debug)]
+enum CallRegionKind {
+    // this is just an intersection of local regions.
+    Intersection(HashSet<Region>),
+    // this is equal to a specific region
+    EqR(Region),
+}
+
+struct CallRegions {
+    pub early_regions: Vec<Region>,
+    pub late_regions: Vec<Region>,
+    pub classification: HashMap<Region, CallRegionKind>,
+}
+
+/// A scope of trait attributes mapping to Coq names to be used in a function's spec
+struct TraitSpecNameScope {
+    attrs: HashMap<String, String>,
+}
+
+/// When translating a function spec where attributes of a trait are in scope,
+/// we create a wrapper to lookup references to the trait's attributes when parsing function specs.
+struct FunctionSpecScope<'a, T> {
+    generics: &'a T,
+    attrs: &'a TraitSpecNameScope,
+}
+impl<'a, T: ParamLookup> ParamLookup for FunctionSpecScope<'a, T> {
+    fn lookup_ty_param(&self, path: &[&str]) -> Option<&radium::LiteralTyParam> {
+        self.generics.lookup_ty_param(path)
+    }
+
+    fn lookup_lft(&self, lft: &str) -> Option<&radium::Lft> {
+        self.generics.lookup_lft(lft)
+    }
+
+    fn lookup_literal(&self, path: &[&str]) -> Option<&str> {
+        if path.len() == 1 {
+            if let Some(lit) = self.attrs.attrs.get(path[0]) {
+                return Some(lit);
+            }
+        }
+        self.generics.lookup_literal(path)
+    }
+}

rr_frontend/translation/src/consts.rs

+// © 2023, The RefinedRust Developers and Contributors
+//
+// This Source Code Form is subject to the terms of the BSD-3-clause License.
+// If a copy of the BSD-3-clause license was not distributed with this
+// file, You can obtain one at https://opensource.org/license/bsd-3-clause/.
+
+use std::collections::{btree_map, BTreeMap, HashMap, HashSet};
+
+use rr_rustc_interface::hir::def_id::DefId;
+
+use crate::base::TranslationError;
+
+/// Scope of consts that are available
+pub struct Scope<'def> {
+    // statics are explicitly declared
+    statics: HashMap<DefId, radium::StaticMeta<'def>>,
+    // const places are constants that lie in a static memory segment because they are referred to
+    // by-ref
+    const_places: HashMap<DefId, radium::ConstPlaceMeta<'def>>,
+}
+
+impl<'def> Scope<'def> {
+    /// Create a new const scope.
+    pub fn empty() -> Self {
+        Self {
+            statics: HashMap::new(),
+            const_places: HashMap::new(),
+        }
+    }
+
+    /// Register a static
+    pub fn register_static(&mut self, did: DefId, meta: radium::StaticMeta<'def>) {
+        self.statics.insert(did, meta);
+    }
+
+    /// Register a const place
+    pub fn register_const_place(&mut self, did: DefId, meta: radium::ConstPlaceMeta<'def>) {
+        self.const_places.insert(did, meta);
+    }
+
+    pub fn get_static<'tcx>(&self, did: DefId) -> Result<&radium::StaticMeta<'def>, TranslationError<'tcx>> {
+        self.statics.get(&did).ok_or_else(|| TranslationError::UnknownError(format!(
+            "Did not find a registered static for did {did:?}; registered: {:?}",
+            self.statics
+        )))
+    }
+}

rr_frontend/translation/src/lib.rs

@@ -11,12 +11,14 @@
 mod arg_folder;
 mod attrs;
 mod base;
+mod body;
 mod checked_op_analysis;
+mod consts;
 mod data;
 pub mod environment;
 mod force_matches_macro;
-mod function_body;
 mod inclusion_tracker;
+mod procedures;
 mod regions;
 mod search;
 mod shims;
@@ -39,8 +41,9 @@ use rr_rustc_interface::{ast, hir, span};
 use topological_sort::TopologicalSort;
 use typed_arena::Arena;
 
+use crate::body::signature;
 use crate::environment::Environment;
-use crate::function_body::{ConstScope, FunctionTranslator, ProcedureMode, ProcedureScope};
+use crate::procedures::{Mode, Scope};
 use crate::shims::registry as shim_registry;
 use crate::spec_parsers::const_attr_parser::{ConstAttrParser, VerboseConstAttrParser};
 use crate::spec_parsers::crate_attr_parser::{CrateAttrParser, VerboseCrateAttrParser};
@@ -80,8 +83,8 @@ fn order_adt_defs(deps: &HashMap<DefId, HashSet<DefId>>) -> Vec<DefId> {
 
 pub struct VerificationCtxt<'tcx, 'rcx> {
     env: &'rcx Environment<'tcx>,
-    procedure_registry: ProcedureScope<'rcx>,
-    const_registry: ConstScope<'rcx>,
+    procedure_registry: procedures::Scope<'rcx>,
+    const_registry: consts::Scope<'rcx>,
     type_translator: &'rcx types::TX<'rcx, 'tcx>,
     trait_registry: &'rcx registry::TR<'tcx, 'rcx>,
     functions: &'rcx [LocalDefId],
@@ -113,7 +116,10 @@ impl<'tcx, 'rcx> VerificationCtxt<'tcx, 'rcx> {
             return None;
         };
 
-        if mode != ProcedureMode::Prove && mode != ProcedureMode::OnlySpec && mode != ProcedureMode::TrustMe {
+        if mode != procedures::Mode::Prove
+            && mode != procedures::Mode::OnlySpec
+            && mode != procedures::Mode::TrustMe
+        {
             return None;
         }
 
@@ -149,7 +155,10 @@ impl<'tcx, 'rcx> VerificationCtxt<'tcx, 'rcx> {
             return None;
         };
 
-        if mode != ProcedureMode::Prove && mode != ProcedureMode::OnlySpec && mode != ProcedureMode::TrustMe {
+        if mode != procedures::Mode::Prove
+            && mode != procedures::Mode::OnlySpec
+            && mode != procedures::Mode::TrustMe
+        {
             trace!("leave make_shim_trait_method_entry (failed)");
             return None;
         }
@@ -924,11 +933,8 @@ fn register_shims<'tcx>(vcx: &mut VerificationCtxt<'tcx, '_>) -> Result<(), base
             Some(did) => {
                 // register as usual in the procedure registry
                 info!("registering shim for {:?}", shim.path);
-                let meta = function_body::ProcedureMeta::new(
-                    shim.spec_name.clone(),
-                    shim.name.clone(),
-                    function_body::ProcedureMode::Shim,
-                );
+                let meta =
+                    procedures::Meta::new(shim.spec_name.clone(), shim.name.clone(), procedures::Mode::Shim);
                 vcx.procedure_registry.register_function(did, meta)?;
             },
             _ => {
@@ -1031,11 +1037,7 @@ fn register_shims<'tcx>(vcx: &mut VerificationCtxt<'tcx, '_>) -> Result<(), base
                     shim.trait_path, method_name, for_type, method_did
                 );
 
-                let meta = function_body::ProcedureMeta::new(
-                    spec_name.clone(),
-                    name.clone(),
-                    function_body::ProcedureMode::Shim,
-                );
+                let meta = procedures::Meta::new(spec_name.clone(), name.clone(), procedures::Mode::Shim);
 
                 vcx.procedure_registry.register_function(method_did, meta)?;
             }
@@ -1051,31 +1053,28 @@ fn register_shims<'tcx>(vcx: &mut VerificationCtxt<'tcx, '_>) -> Result<(), base
     Ok(())
 }
 
-fn get_most_restrictive_function_mode(
-    vcx: &VerificationCtxt<'_, '_>,
-    did: DefId,
-) -> function_body::ProcedureMode {
+fn get_most_restrictive_function_mode(vcx: &VerificationCtxt<'_, '_>, did: DefId) -> procedures::Mode {
     let attrs = vcx.env.get_attributes_of_function(did, &propagate_method_attr_from_impl);
 
     // check if this is a purely spec function; if so, skip.
     if attrs::has_tool_attr_filtered(attrs.as_slice(), "shim") {
-        return function_body::ProcedureMode::Shim;
+        return procedures::Mode::Shim;
     }
 
     if attrs::has_tool_attr_filtered(attrs.as_slice(), "trust_me") {
-        return function_body::ProcedureMode::TrustMe;
+        return procedures::Mode::TrustMe;
     }
 
     if attrs::has_tool_attr_filtered(attrs.as_slice(), "only_spec") {
-        return function_body::ProcedureMode::OnlySpec;
+        return procedures::Mode::OnlySpec;
     }
 
     if attrs::has_tool_attr_filtered(attrs.as_slice(), "ignore") {
         info!("Function {:?} will be ignored", did);
-        return function_body::ProcedureMode::Ignore;
+        return procedures::Mode::Ignore;
     }
 
-    function_body::ProcedureMode::Prove
+    procedures::Mode::Prove
 }
 
 /// Register functions of the crate in the procedure registry.
@@ -1085,7 +1084,7 @@ fn register_functions<'tcx>(
     for &f in vcx.functions {
         let mut mode = get_most_restrictive_function_mode(vcx, f.to_def_id());
 
-        if mode == function_body::ProcedureMode::Shim {
+        if mode == procedures::Mode::Shim {
             // TODO better error message
             let attrs = vcx.env.get_attributes(f.to_def_id());
             let v = attrs::filter_for_tool(attrs);
@@ -1097,29 +1096,25 @@ fn register_functions<'tcx>(
                 annot.spec_name,
                 annot.code_name
             );
-            let meta = function_body::ProcedureMeta::new(
-                annot.spec_name,
-                annot.code_name,
-                function_body::ProcedureMode::Shim,
-            );
+            let meta = procedures::Meta::new(annot.spec_name, annot.code_name, procedures::Mode::Shim);
             vcx.procedure_registry.register_function(f.to_def_id(), meta)?;
 
             continue;
         }
 
-        if mode == function_body::ProcedureMode::Prove && let Some(impl_did) = vcx.env.tcx().impl_of_method(f.to_def_id()) {
+        if mode == procedures::Mode::Prove && let Some(impl_did) = vcx.env.tcx().impl_of_method(f.to_def_id()) {
             mode = get_most_restrictive_function_mode(vcx, impl_did);
         }
 
-        if mode == function_body::ProcedureMode::Shim {
+        if mode == procedures::Mode::Shim {
             warn!("Nonsensical shim attribute on impl; ignoring");
-            mode = function_body::ProcedureMode::Prove;
+            mode = procedures::Mode::Prove;
         }
 
         let fname = base::strip_coq_ident(&vcx.env.get_item_name(f.to_def_id()));
         let spec_name = format!("type_of_{}", fname);
 
-        let meta = function_body::ProcedureMeta::new(spec_name, fname, mode);
+        let meta = procedures::Meta::new(spec_name, fname, mode);
 
         vcx.procedure_registry.register_function(f.to_def_id(), meta)?;
     }
@@ -1150,7 +1145,7 @@ fn translate_functions<'rcx, 'tcx>(vcx: &mut VerificationCtxt<'tcx, 'rcx>) {
         let ty = ty.instantiate_identity();
 
         let translator = match ty.kind() {
-            ty::TyKind::FnDef(_def, _args) => FunctionTranslator::new(
+            ty::TyKind::FnDef(_def, _args) => signature::TX::new(
                 vcx.env,
                 &meta,
                 proc,
@@ -1160,7 +1155,7 @@ fn translate_functions<'rcx, 'tcx>(vcx: &mut VerificationCtxt<'tcx, 'rcx>) {
                 &vcx.procedure_registry,
                 &vcx.const_registry,
             ),
-            ty::TyKind::Closure(_, _) => FunctionTranslator::new_closure(
+            ty::TyKind::Closure(_, _) => signature::TX::new_closure(
                 vcx.env,
                 &meta,
                 proc,
@@ -1175,7 +1170,7 @@ fn translate_functions<'rcx, 'tcx>(vcx: &mut VerificationCtxt<'tcx, 'rcx>) {
 
         if mode.is_only_spec() {
             // Only generate a spec
-            match translator.and_then(FunctionTranslator::generate_spec) {
+            match translator.and_then(signature::TX::generate_spec) {
                 Ok(spec) => {
                     println!("Successfully generated spec for {}", fname);
                     let spec_ref = vcx.fn_arena.alloc(spec);
@@ -1533,7 +1528,7 @@ where
         &trait_use_arena,
         &fn_spec_arena,
     );
-    let procedure_registry = ProcedureScope::new();
+    let procedure_registry = procedures::Scope::new();
     let shim_string_arena = Arena::new();
     let mut shim_registry = shim_registry::SR::empty(&shim_string_arena);
 
@@ -1574,7 +1569,7 @@ where
         coq_path_prefix: path_prefix,
         shim_registry,
         dune_package: package,
-        const_registry: ConstScope::empty(),
+        const_registry: consts::Scope::empty(),
         trait_impls: HashMap::new(),
         fn_arena: &fn_spec_arena,
     };

rr_frontend/translation/src/procedures.rs

+// © 2023, The RefinedRust Developers and Contributors
+//
+// This Source Code Form is subject to the terms of the BSD-3-clause License.
+// If a copy of the BSD-3-clause license was not distributed with this
+// file, You can obtain one at https://opensource.org/license/bsd-3-clause/.
+
+use std::collections::{btree_map, BTreeMap, HashMap, HashSet};
+
+use rr_rustc_interface::hir::def_id::DefId;
+
+use crate::base::TranslationError;
+
+#[derive(Copy, Clone, Eq, PartialEq, Debug)]
+pub enum Mode {
+    Prove,
+    OnlySpec,
+    TrustMe,
+    Shim,
+    Ignore,
+}
+
+impl Mode {
+    pub fn is_prove(self) -> bool {
+        self == Self::Prove
+    }
+
+    pub fn is_only_spec(self) -> bool {
+        self == Self::OnlySpec
+    }
+
+    pub fn is_trust_me(self) -> bool {
+        self == Self::TrustMe
+    }
+
+    pub fn is_shim(self) -> bool {
+        self == Self::Shim
+    }
+
+    pub fn is_ignore(self) -> bool {
+        self == Self::Ignore
+    }
+
+    pub fn needs_proof(self) -> bool {
+        self == Self::Prove
+    }
+
+    pub fn needs_def(self) -> bool {
+        self == Self::Prove || self == Self::TrustMe
+    }
+
+    pub fn needs_spec(self) -> bool {
+        self == Self::Prove || self == Self::TrustMe || self == Self::OnlySpec
+    }
+}
+
+#[derive(Clone)]
+pub struct Meta {
+    spec_name: String,
+    name: String,
+    mode: Mode,
+}
+
+impl Meta {
+    pub const fn new(spec_name: String, name: String, mode: Mode) -> Self {
+        Self {
+            spec_name,
+            name,
+            mode,
+        }
+    }
+
+    pub fn get_spec_name(&self) -> &str {
+        &self.spec_name
+    }
+
+    pub fn get_name(&self) -> &str {
+        &self.name
+    }
+
+    pub const fn get_mode(&self) -> Mode {
+        self.mode
+    }
+}
+
+/**
+ * Tracks the functions we translated and the Coq names they are available under.
+ * To account for dependencies between functions, we may register translated names before we have
+ * actually translated the function.
+ */
+pub struct Scope<'def> {
+    /// maps the defid to (code_name, spec_name, name)
+    name_map: BTreeMap<DefId, Meta>,
+    /// track the actually translated functions
+    translated_functions: BTreeMap<DefId, radium::Function<'def>>,
+    /// track the functions with just a specification (rr::only_spec)
+    specced_functions: BTreeMap<DefId, &'def radium::FunctionSpec<'def, radium::InnerFunctionSpec<'def>>>,
+}
+
+impl<'def> Scope<'def> {
+    pub fn new() -> Self {
+        Self {
+            name_map: BTreeMap::new(),
+            translated_functions: BTreeMap::new(),
+            specced_functions: BTreeMap::new(),
+        }
+    }
+
+    /// Lookup the meta information of a function.
+    pub fn lookup_function(&self, did: DefId) -> Option<Meta> {
+        self.name_map.get(&did).cloned()
+    }
+
+    /// Lookup a translated function spec
+    pub fn lookup_function_spec(
+        &self,
+        did: DefId,
+    ) -> Option<&'def radium::FunctionSpec<'def, radium::InnerFunctionSpec<'def>>> {
+        if let Some(translated_fn) = self.translated_functions.get(&did) {
+            Some(translated_fn.spec)
+        } else if let Some(translated_spec) = self.specced_functions.get(&did) {
+            Some(translated_spec)
+        } else {
+            None
+        }
+    }
+
+    /// Lookup the Coq spec name for a function.
+    pub fn lookup_function_spec_name(&self, did: DefId) -> Option<&str> {
+        self.name_map.get(&did).map(Meta::get_spec_name)
+    }
+
+    /// Lookup the name for a function.
+    pub fn lookup_function_mangled_name(&self, did: DefId) -> Option<&str> {
+        self.name_map.get(&did).map(Meta::get_name)
+    }
+
+    /// Lookup the mode for a function.
+    pub fn lookup_function_mode(&self, did: DefId) -> Option<Mode> {
+        self.name_map.get(&did).map(Meta::get_mode)
+    }
+
+    /// Register a function.
+    pub fn register_function<'tcx>(&mut self, did: DefId, meta: Meta) -> Result<(), TranslationError<'tcx>> {
+        if self.name_map.insert(did, meta).is_some() {
+            Err(TranslationError::ProcedureRegistry(format!(
+                "function for defid {:?} has already been registered",
+                did
+            )))
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Provide the code for a translated function.
+    pub fn provide_translated_function(&mut self, did: DefId, trf: radium::Function<'def>) {
+        let meta = &self.name_map[&did];
+        assert!(meta.get_mode().needs_def());
+        assert!(self.translated_functions.insert(did, trf).is_none());
+    }
+
+    /// Provide the specification for an `only_spec` function.
+    pub fn provide_specced_function(
+        &mut self,
+        did: DefId,
+        spec: &'def radium::FunctionSpec<'def, radium::InnerFunctionSpec<'def>>,
+    ) {
+        let meta = &self.name_map[&did];
+        assert!(meta.get_mode().is_only_spec());
+        assert!(self.specced_functions.insert(did, spec).is_none());
+    }
+
+    /// Iterate over the functions we have generated code for.
+    pub fn iter_code(&self) -> btree_map::Iter<'_, DefId, radium::Function<'def>> {
+        self.translated_functions.iter()
+    }
+
+    /// Iterate over the functions we have generated only specs for.
+    pub fn iter_only_spec(
+        &self,
+    ) -> btree_map::Iter<'_, DefId, &'def radium::FunctionSpec<'def, radium::InnerFunctionSpec<'def>>> {
+        self.specced_functions.iter()
+    }
+}

rr_frontend/translation/src/regions/mod.rs

@@ -9,13 +9,15 @@
 use std::collections::{BTreeMap, HashMap, HashSet};
 
 use derive_more::{Constructor, Debug};
+use log::{info, warn};
 use rr_rustc_interface::hir::def_id::DefId;
 use rr_rustc_interface::middle::mir::tcx::PlaceTy;
 use rr_rustc_interface::middle::ty;
 use rr_rustc_interface::middle::ty::{Ty, TyCtxt, TyKind, TypeFoldable};
 use ty::TypeSuperFoldable;
 
-use crate::base::*;
+use crate::arg_folder::ty_instantiate;
+use crate::base::{self, Region};
 use crate::environment::polonius_info::{self, PoloniusInfo};
 use crate::environment::Environment;
 
@@ -72,7 +74,7 @@ pub fn find_region_variables_of_place_type<'tcx>(
     collector.get_regions()
 }
 
-/// Data structure that maps early and late region indices to Polonius regions.
+/// Data structure that maps early and late region indices inside functions to Polonius regions.
 #[derive(Constructor, Clone, Debug, Default)]
 pub struct EarlyLateRegionMap {
     // maps indices of early and late regions to Polonius region ids
@@ -86,6 +88,26 @@ pub struct EarlyLateRegionMap {
     pub lft_names: HashMap<String, ty::RegionVid>,
 }
 impl EarlyLateRegionMap {
+    /// Lookup a Polonius region with a given kind.
+    pub fn lookup_region_with_kind(
+        &self,
+        k: polonius_info::UniversalRegionKind,
+        r: Region,
+    ) -> Option<radium::UniversalLft> {
+        match k {
+            polonius_info::UniversalRegionKind::Function => Some(radium::UniversalLft::Function),
+            polonius_info::UniversalRegionKind::Static => Some(radium::UniversalLft::Static),
+
+            polonius_info::UniversalRegionKind::Local => {
+                self.lookup_region(r).map(|x| radium::UniversalLft::Local(x.to_owned()))
+            },
+
+            polonius_info::UniversalRegionKind::External => {
+                self.lookup_region(r).map(|x| radium::UniversalLft::External(x.to_owned()))
+            },
+        }
+    }
+
     pub fn lookup_region(&self, region: ty::RegionVid) -> Option<&radium::Lft> {
         self.region_names.get(&region)
     }
@@ -125,3 +147,115 @@ pub fn format_atomic_region_direct(
         },
     }
 }
+
+/// Process the signature of a function by instantiating the region variables with their
+/// Polonius variables.
+/// Returns the argument and return type signature instantiated in this way.
+/// Moreover, returns a `EarlyLateRegionMap` that contains the mapping of indices to Polonius
+/// region variables.
+pub fn replace_fnsig_args_with_polonius_vars<'tcx>(
+    env: &Environment<'tcx>,
+    params: &[ty::GenericArg<'tcx>],
+    sig: ty::Binder<'tcx, ty::FnSig<'tcx>>,
+) -> (Vec<ty::Ty<'tcx>>, ty::Ty<'tcx>, EarlyLateRegionMap) {
+    // a mapping of Polonius region IDs to names
+    let mut universal_lifetimes = BTreeMap::new();
+    let mut lifetime_names = HashMap::new();
+
+    let mut region_substitution_early: Vec<Option<ty::RegionVid>> = Vec::new();
+
+    // we create a substitution that replaces early bound regions with their Polonius
+    // region variables
+    let mut subst_early_bounds: Vec<ty::GenericArg<'tcx>> = Vec::new();
+    let mut num_early_bounds = 0;
+    for a in params {
+        if let ty::GenericArgKind::Lifetime(r) = a.unpack() {
+            // skip over 0 = static
+            let next_id = ty::RegionVid::from_u32(num_early_bounds + 1);
+            let revar = ty::Region::new_var(env.tcx(), next_id);
+            num_early_bounds += 1;
+            subst_early_bounds.push(ty::GenericArg::from(revar));
+
+            region_substitution_early.push(Some(next_id));
+
+            match *r {
+                ty::RegionKind::ReEarlyBound(r) => {
+                    let name = base::strip_coq_ident(r.name.as_str());
+                    universal_lifetimes.insert(next_id, format!("ulft_{}", name));
+                    lifetime_names.insert(name, next_id);
+                },
+                _ => {
+                    universal_lifetimes.insert(next_id, format!("ulft_{}", num_early_bounds));
+                },
+            }
+        } else {
+            subst_early_bounds.push(*a);
+            region_substitution_early.push(None);
+        }
+    }
+    let subst_early_bounds = env.tcx().mk_args(&subst_early_bounds);
+
+    info!("Computed early region map {region_substitution_early:?}");
+
+    // add names for late bound region variables
+    let mut num_late_bounds = 0;
+    let mut region_substitution_late = Vec::new();
+    for b in sig.bound_vars() {
+        let next_id = ty::RegionVid::from_u32(num_early_bounds + num_late_bounds + 1);
+
+        let ty::BoundVariableKind::Region(r) = b else {
+            continue;
+        };
+
+        region_substitution_late.push(next_id);
+
+        match r {
+            ty::BoundRegionKind::BrNamed(_, sym) => {
+                let mut region_name = base::strip_coq_ident(sym.as_str());
+                if region_name == "_" {
+                    region_name = next_id.as_usize().to_string();
+                    universal_lifetimes.insert(next_id, format!("ulft_{}", region_name));
+                } else {
+                    universal_lifetimes.insert(next_id, format!("ulft_{}", region_name));
+                    lifetime_names.insert(region_name, next_id);
+                }
+            },
+            ty::BoundRegionKind::BrAnon(_) => {
+                universal_lifetimes.insert(next_id, format!("ulft_{}", next_id.as_usize()));
+            },
+            _ => (),
+        }
+
+        num_late_bounds += 1;
+    }
+
+    // replace late-bound region variables by re-enumerating them in the same way as the MIR
+    // type checker does (that this happens in the same way is important to make the names
+    // line up!)
+    let mut next_index = num_early_bounds + 1; // skip over one additional due to static
+    let mut folder = |_| {
+        let cur_index = next_index;
+        next_index += 1;
+        ty::Region::new_var(env.tcx(), ty::RegionVid::from_u32(cur_index))
+    };
+    let (late_sig, _late_region_map) = env.tcx().replace_late_bound_regions(sig, &mut folder);
+
+    // replace early bound variables
+    let inputs: Vec<_> = late_sig
+        .inputs()
+        .iter()
+        .map(|ty| ty_instantiate(*ty, env.tcx(), subst_early_bounds))
+        .collect();
+
+    let output = ty_instantiate(late_sig.output(), env.tcx(), subst_early_bounds);
+
+    info!("Computed late region map {region_substitution_late:?}");
+
+    let region_map = EarlyLateRegionMap::new(
+        region_substitution_early,
+        region_substitution_late,
+        universal_lifetimes,
+        lifetime_names,
+    );
+    (inputs, output, region_map)
+}

rr_frontend/translation/src/shims/flat.rs

@@ -7,9 +7,9 @@
 //! Provides a flat representation of types that is stable across compilations.
 
 use log::{info, trace};
+use rr_rustc_interface::hir;
 use rr_rustc_interface::hir::def_id::DefId;
 use rr_rustc_interface::middle::ty::{self, TyCtxt};
-use rr_rustc_interface::hir;
 use serde::{Deserialize, Serialize};
 
 use crate::spec_parsers::get_export_as_attr;

rr_frontend/translation/src/spec_parsers/verbose_function_spec_parser.rs

@@ -242,6 +242,7 @@ pub struct VerboseFunctionSpecParser<'a, 'def, F, T> {
     got_ret: bool,
 }
 
+/// Extra requirements of a function.
 #[derive(Default)]
 pub struct FunctionRequirements {
     /// additional late coq parameters

rr_frontend/translation/src/traits/registry.rs

@@ -16,8 +16,8 @@ use traits::{resolution, Error, TraitResult};
 use typed_arena::Arena;
 
 use crate::base::TranslationError;
+use crate::body::signature;
 use crate::environment::Environment;
-use crate::function_body::FunctionTranslator;
 use crate::spec_parsers::propagate_method_attr_from_impl;
 use crate::spec_parsers::trait_attr_parser::{TraitAttrParser, VerboseTraitAttrParser};
 use crate::spec_parsers::trait_impl_attr_parser::{TraitImplAttrParser, VerboseTraitImplAttrParser};
@@ -223,7 +223,7 @@ impl<'tcx, 'def> TR<'tcx, 'def> {
                 let spec_name = format!("{name}_base_spec");
 
                 // get spec
-                let spec = FunctionTranslator::spec_for_trait_method(
+                let spec = signature::TX::spec_for_trait_method(
                     self.env,
                     c.def_id,
                     &name,