[move-prover] better architecture for simplification pipeline

This creates a new trait, `SpecRewriter` that handles the
transformations on specs.

The `SpecRewriter` has three functions to be implemented by each
individual rewriter:
- `rewrite_module_spec`
- `rewrite_function_spec`
- `rewrite_inline_spec`

Multiple rewriters that implement the `SpecRewriter` trait can be
composed into a new rewriter that pipelines them. This logic is captured
in `SpecRewriterPipeline`, which also implements `SpecRewriter`.

A `SpecRewriter` has an option to write-back the newly generated spec
into a `&mut GlobalEnv` with the `override_with_rewrite` function. But
this is optional and only available in cases where a `&mut GlobalEnv`
can be obtained (e.g., during model building or right after the model
is built). If writing-back is not needed, the three `rewrite_*` funciton
should be sufficient.

Closes: aptos-labs#9821

language/move-model/src/lib.rs

@@ -39,6 +39,7 @@ use crate::{
     builder::model_builder::ModelBuilder,
     model::{FunId, FunctionData, GlobalEnv, Loc, ModuleData, ModuleId, StructId},
     options::ModelBuilderOptions,
+    simplifier::{SpecRewriter, SpecRewriterPipeline},
 };
 
 pub mod ast;
@@ -581,10 +582,10 @@ fn run_spec_simplifier(env: &mut GlobalEnv) {
     let options = env
         .get_extension::<ModelBuilderOptions>()
         .expect("options for model builder");
-    for pass in &options.simplification_pipeline {
-        pass.run(env)
-            .unwrap_or_else(|e| panic!("Failed to run simplification pass {}: {}", pass, e))
-    }
+    let mut rewriter = SpecRewriterPipeline::new(&options.simplification_pipeline);
+    rewriter
+        .override_with_rewrite(env)
+        .unwrap_or_else(|e| panic!("Failed to run spec simplification: {}", e))
 }
 
 // =================================================================================================

language/move-model/src/model.rs

@@ -70,6 +70,7 @@ use crate::{
 };
 
 // import and re-expose symbols
+use move_binary_format::file_format::CodeOffset;
 pub use move_binary_format::file_format::{AbilitySet, Visibility as FunctionVisibility};
 
 // =================================================================================================
@@ -1518,6 +1519,19 @@ impl GlobalEnv {
         total
     }
 
+    /// Override the specification for a given module
+    pub fn override_module_spec(&mut self, mid: ModuleId, spec: Spec) {
+        let module_data = self
+            .module_data
+            .iter_mut()
+            .filter(|m| m.id == mid)
+            .exactly_one()
+            .unwrap_or_else(|_| {
+                panic!("Expect one and only one module for {:?}", mid);
+            });
+        module_data.module_spec = spec;
+    }
+
     /// Override the specification for a given function
     pub fn override_function_spec(&mut self, fid: QualifiedId<FunId>, spec: Spec) {
         let func_data = self
@@ -1537,6 +1551,31 @@ impl GlobalEnv {
             });
         func_data.spec = spec;
     }
+
+    /// Override the specification for a given code location
+    pub fn override_inline_spec(
+        &mut self,
+        fid: QualifiedId<FunId>,
+        code_offset: CodeOffset,
+        spec: Spec,
+    ) {
+        let func_data = self
+            .module_data
+            .iter_mut()
+            .filter(|m| m.id == fid.module_id)
+            .map(|m| {
+                m.function_data
+                    .iter_mut()
+                    .filter(|(k, _)| **k == fid.id)
+                    .map(|(_, v)| v)
+            })
+            .flatten()
+            .exactly_one()
+            .unwrap_or_else(|_| {
+                panic!("Expect one and only one function for {:?}", fid);
+            });
+        func_data.spec.on_impl.insert(code_offset, spec);
+    }
 }
 
 impl Default for GlobalEnv {

language/move-model/src/simplifier/mod.rs

@@ -5,10 +5,19 @@ use anyhow::Result;
 use serde::{Deserialize, Serialize};
 use std::{fmt, str::FromStr};
 
-use crate::model::GlobalEnv;
+use move_binary_format::file_format::CodeOffset;
 
+use crate::{
+    ast::Spec,
+    model::{FunId, GlobalEnv, ModuleId, QualifiedId},
+};
+
+mod pass;
 mod pass_inline;
 
+pub use pass::SpecRewriter;
+use pass_inline::SpecPassInline;
+
 /// Available simplifications passes to run after tbe model is built
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub enum SimplificationPass {
@@ -35,10 +44,84 @@ impl fmt::Display for SimplificationPass {
     }
 }
 
-impl SimplificationPass {
-    pub fn run(&self, env: &mut GlobalEnv) -> Result<()> {
-        match self {
-            Self::Inline => pass_inline::run_pass_inline(env),
+/// A rewriter pipeline that is composed of a chain of spec rewriters. Note that this composite
+/// rewriter is also a spec rewriter.
+pub struct SpecRewriterPipeline {
+    rewriters: Vec<Box<dyn SpecRewriter>>,
+}
+
+impl SpecRewriterPipeline {
+    /// Construct a pipeline rewriter by a list of passes
+    pub fn new(pipeline: &[SimplificationPass]) -> Self {
+        let mut result = Self { rewriters: vec![] };
+        for entry in pipeline {
+            match entry {
+                SimplificationPass::Inline => {
+                    result.rewriters.push(Box::new(SpecPassInline::default()))
+                }
+            }
+        }
+        result
+    }
+}
+
+impl SpecRewriter for SpecRewriterPipeline {
+    fn rewrite_module_spec(
+        &mut self,
+        env: &GlobalEnv,
+        module_id: ModuleId,
+        spec: &Spec,
+    ) -> Result<Option<Spec>> {
+        let mut current_spec = None;
+        for rewriter in self.rewriters.iter_mut() {
+            if let Some(new_spec) = rewriter.rewrite_module_spec(
+                env,
+                module_id,
+                current_spec.as_ref().unwrap_or(spec),
+            )? {
+                current_spec = Some(new_spec);
+            }
+        }
+        Ok(current_spec)
+    }
+
+    fn rewrite_function_spec(
+        &mut self,
+        env: &GlobalEnv,
+        fun_id: QualifiedId<FunId>,
+        spec: &Spec,
+    ) -> Result<Option<Spec>> {
+        let mut current_spec = None;
+        for rewriter in self.rewriters.iter_mut() {
+            if let Some(new_spec) = rewriter.rewrite_function_spec(
+                env,
+                fun_id,
+                current_spec.as_ref().unwrap_or(spec),
+            )? {
+                current_spec = Some(new_spec);
+            }
+        }
+        Ok(current_spec)
+    }
+
+    fn rewrite_inline_spec(
+        &mut self,
+        env: &GlobalEnv,
+        fun_id: QualifiedId<FunId>,
+        code_offset: CodeOffset,
+        spec: &Spec,
+    ) -> Result<Option<Spec>> {
+        let mut current_spec = None;
+        for rewriter in self.rewriters.iter_mut() {
+            if let Some(new_spec) = rewriter.rewrite_inline_spec(
+                env,
+                fun_id,
+                code_offset,
+                current_spec.as_ref().unwrap_or(spec),
+            )? {
+                current_spec = Some(new_spec);
+            }
         }
+        Ok(current_spec)
     }
 }

language/move-model/src/simplifier/pass.rs

@@ -0,0 +1,98 @@
+// Copyright (c) The Diem Core Contributors
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::Result;
+use std::collections::BTreeMap;
+
+use move_binary_format::file_format::CodeOffset;
+
+use crate::{
+    ast::Spec,
+    model::{FunId, GlobalEnv, ModuleId, QualifiedId},
+};
+
+/// A generic trait for rewriting the specifications in the `GlobalEnv`. A rewriter is expected to
+/// implement at least one `rewrite_*` function, depending on which type(s) of specs the rewriter
+/// targets. All the `rewrite_*` function should follow the convention on return value:
+/// - `Ok(None)`           --> nothing to rewrite on this spec
+/// - `Ok(Some(new_spec))` --> the spec is re-written and the `new_spec` is the output
+/// - `Err(..)`            --> something wrong (invariant violation) happened in the rewriting
+pub trait SpecRewriter {
+    /// Rewrite a module-level specification
+    fn rewrite_module_spec(
+        &mut self,
+        _env: &GlobalEnv,
+        _module_id: ModuleId,
+        _spec: &Spec,
+    ) -> Result<Option<Spec>> {
+        Ok(None)
+    }
+
+    /// Rewrite a function-level specification
+    fn rewrite_function_spec(
+        &mut self,
+        _env: &GlobalEnv,
+        _fun_id: QualifiedId<FunId>,
+        _spec: &Spec,
+    ) -> Result<Option<Spec>> {
+        Ok(None)
+    }
+
+    /// Rewrite a code-level specification
+    fn rewrite_inline_spec(
+        &mut self,
+        _env: &GlobalEnv,
+        _fun_id: QualifiedId<FunId>,
+        _code_offset: CodeOffset,
+        _spec: &Spec,
+    ) -> Result<Option<Spec>> {
+        Ok(None)
+    }
+
+    /// Iterate over the specs in the `GlobalEnv`, rewrite each spec, and apply changes back to the
+    /// `GlobalEnv`.
+    fn override_with_rewrite(&mut self, env: &mut GlobalEnv) -> Result<()> {
+        // convert all module specs found in the model
+        let mut new_specs = BTreeMap::new();
+        for menv in env.get_modules() {
+            let mid = menv.get_id();
+            if let Some(new_spec) = self.rewrite_module_spec(env, mid, menv.get_spec())? {
+                new_specs.insert(mid, new_spec);
+            }
+        }
+        for (mid, spec) in new_specs {
+            env.override_module_spec(mid, spec);
+        }
+
+        // convert all functional specs found in the model
+        let mut new_specs = BTreeMap::new();
+        for menv in env.get_modules() {
+            for fenv in menv.get_functions() {
+                let fid = fenv.get_qualified_id();
+                if let Some(new_spec) = self.rewrite_function_spec(env, fid, fenv.get_spec())? {
+                    new_specs.insert(fid, new_spec);
+                }
+            }
+        }
+        for (fid, spec) in new_specs {
+            env.override_function_spec(fid, spec);
+        }
+
+        // convert all code-level specs found in the model
+        let mut new_specs = BTreeMap::new();
+        for menv in env.get_modules() {
+            for fenv in menv.get_functions() {
+                let fid = fenv.get_qualified_id();
+                for (offset, spec) in &fenv.get_spec().on_impl {
+                    if let Some(new_spec) = self.rewrite_inline_spec(env, fid, *offset, spec)? {
+                        new_specs.insert((fid, *offset), new_spec);
+                    }
+                }
+            }
+        }
+        for ((fid, offset), spec) in new_specs {
+            env.override_inline_spec(fid, offset, spec);
+        }
+        Ok(())
+    }
+}