Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement MastForest merging #1534

Merged
merged 45 commits into from
Oct 28, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
45 commits
Select commit Hold shift + click to select a range
82b90c1
feat(core): Move `EqHash` to core and rename
PhilippGackstatter Oct 14, 2024
2f4529b
feat(core): Implement Mast Forest merging
PhilippGackstatter Oct 15, 2024
c431f64
feat(core): Address clippy lints
PhilippGackstatter Oct 16, 2024
ff15c84
feat(core): Optimize id mappings with `ForestIdMap`
PhilippGackstatter Oct 16, 2024
0a1310e
feat(core): Document code and improve code quality
PhilippGackstatter Oct 16, 2024
e83b573
feat(core): Move dfs iterator to its own module
PhilippGackstatter Oct 16, 2024
d2983d6
feat(core): Add documentation
PhilippGackstatter Oct 16, 2024
00109dd
fix(core): Dfs iterator on empty forest
PhilippGackstatter Oct 16, 2024
bc20663
feat(core): Revert `MastNodeEq` rename
PhilippGackstatter Oct 16, 2024
2892036
chore(core): Simplify and improve structure
PhilippGackstatter Oct 16, 2024
38e048a
chore: Add changelog entry
PhilippGackstatter Oct 16, 2024
cf7bac8
fix(core): Use `roots` of the forest for discovering trees
PhilippGackstatter Oct 18, 2024
0675413
feat(core): Rework API and adapt to changed DFS iteration behavior
PhilippGackstatter Oct 18, 2024
5ec8eee
feat(core): Add `MastForest::iter_nodes`
PhilippGackstatter Oct 18, 2024
375d36c
chore(core): Move tests to separate file
PhilippGackstatter Oct 18, 2024
9625fc0
feat(core): Implement dedicated external node merging
PhilippGackstatter Oct 18, 2024
22efb8c
feat(core): Add test cases for external nodes with decorators
PhilippGackstatter Oct 18, 2024
4a7e28c
fix(core): Implement "drop external node" merging behavior
PhilippGackstatter Oct 18, 2024
5fec6df
feat(core): Deduplicate external nodes with same MAST root
PhilippGackstatter Oct 21, 2024
2887e74
feat(core): Make `merge` impl independent of `&self`
PhilippGackstatter Oct 21, 2024
62a1d5d
chore(core): Update doc comments on iterator
PhilippGackstatter Oct 21, 2024
535c11b
feat(core): Handle invalid decorator indices
PhilippGackstatter Oct 21, 2024
f5e8e97
feat(core): Implement `MastForestRootMap`
PhilippGackstatter Oct 21, 2024
883f224
chore(core): Update docs to new implementation
PhilippGackstatter Oct 21, 2024
c976b4c
feat(core): Add `DecoratorId::new_unsafe` constructor
PhilippGackstatter Oct 21, 2024
72e99e5
chore(core): Make root map pub, fix doc link
PhilippGackstatter Oct 21, 2024
f64b94b
feat(assembly): Add assembler-produced forest merge test
PhilippGackstatter Oct 21, 2024
462ce59
chore(core): Rename `dfs_iterator` mod
PhilippGackstatter Oct 22, 2024
9af5a6c
chore(core): Change `lookup_external_node_by_root` param
PhilippGackstatter Oct 22, 2024
33a4b9c
feat(core): Take ownership of merging forests
PhilippGackstatter Oct 22, 2024
56ae0d9
Revert "feat(core): Take ownership of merging forests"
PhilippGackstatter Oct 23, 2024
a761e9f
feat(core): Implement Multi forest iterator
PhilippGackstatter Oct 23, 2024
5ab8d56
chore: fix lints and errors
bobbinth Oct 24, 2024
d46a073
chore(core): Add missing documentation and cleanup
PhilippGackstatter Oct 24, 2024
644fcb0
chore(core): Use index with fingerprints
PhilippGackstatter Oct 24, 2024
617e514
feat(core): Move Vec inside the `MastForestRootMap`
PhilippGackstatter Oct 24, 2024
7082f55
chore: Address review comments
PhilippGackstatter Oct 24, 2024
d469149
feat(core): Allow unvisited nodes on the iterator stack
PhilippGackstatter Oct 24, 2024
b4c48d5
feat(core): Add test case for duplicate children
PhilippGackstatter Oct 24, 2024
f09d7f9
chore(stdlib): Add mast forest merge test
PhilippGackstatter Oct 24, 2024
f5a7327
chore(core): Rename to `{MastNodeId,DecoratorId}::new_unchecked`
PhilippGackstatter Oct 24, 2024
2fbe528
feat(core): Use a deque instead of a stack in the iterator
PhilippGackstatter Oct 25, 2024
9d5400f
chore(core): Improve docs and comments
PhilippGackstatter Oct 25, 2024
660ad95
fix(core): Use `VecDeque` from `alloc` instead of `std`
PhilippGackstatter Oct 25, 2024
0f76408
chore(stdlib): Don't clone stdlib forest
PhilippGackstatter Oct 25, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
- [BREAKING] The `run` and the `prove` commands in the cli will accept `--trace` flag instead of `--tracing` (#1502)
- Migrated to new padding rule for RPO (#1343).
- Migrated to `miden-crypto` v0.11.0 (#1343).
- Implemented `MastForest` merging (#1534)

#### Fixes

Expand Down
139 changes: 3 additions & 136 deletions assembly/src/assembler/mast_forest_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ use alloc::{
use core::ops::{Index, IndexMut};

use vm_core::{
crypto::hash::{Blake3Digest, Blake3_256, Digest, RpoDigest},
mast::{DecoratorId, MastForest, MastNode, MastNodeId},
crypto::hash::{Blake3Digest, RpoDigest},
mast::{DecoratorId, EqHash, MastForest, MastNode, MastNodeId},
Decorator, DecoratorList, Operation,
};

Expand Down Expand Up @@ -445,115 +445,9 @@ impl MastForestBuilder {
}
}

/// Helpers
impl MastForestBuilder {
fn eq_hash_for_node(&self, node: &MastNode) -> EqHash {
match node {
MastNode::Block(node) => {
let mut bytes_to_hash = Vec::new();

for &(idx, decorator_id) in node.decorators() {
bytes_to_hash.extend(idx.to_le_bytes());
bytes_to_hash.extend(self[decorator_id].eq_hash().as_bytes());
}

// Add any `Assert` or `U32assert2` opcodes present, since these are not included in
// the MAST root.
for (op_idx, op) in node.operations().enumerate() {
if let Operation::U32assert2(inner_value)
| Operation::Assert(inner_value)
| Operation::MpVerify(inner_value) = op
{
let op_idx: u32 = op_idx
.try_into()
.expect("there are more than 2^{32}-1 operations in basic block");

// we include the opcode to differentiate between `Assert` and `U32assert2`
bytes_to_hash.push(op.op_code());
// we include the operation index to distinguish between basic blocks that
// would have the same assert instructions, but in a different order
bytes_to_hash.extend(op_idx.to_le_bytes());
bytes_to_hash.extend(inner_value.to_le_bytes());
}
}

if bytes_to_hash.is_empty() {
EqHash::new(node.digest())
} else {
let decorator_root = Blake3_256::hash(&bytes_to_hash);
EqHash::with_decorator_root(node.digest(), decorator_root)
}
},
MastNode::Join(node) => self.eq_hash_from_parts(
node.before_enter(),
node.after_exit(),
&[node.first(), node.second()],
node.digest(),
),
MastNode::Split(node) => self.eq_hash_from_parts(
node.before_enter(),
node.after_exit(),
&[node.on_true(), node.on_false()],
node.digest(),
),
MastNode::Loop(node) => self.eq_hash_from_parts(
node.before_enter(),
node.after_exit(),
&[node.body()],
node.digest(),
),
MastNode::Call(node) => self.eq_hash_from_parts(
node.before_enter(),
node.after_exit(),
&[node.callee()],
node.digest(),
),
MastNode::Dyn(node) => {
self.eq_hash_from_parts(node.before_enter(), node.after_exit(), &[], node.digest())
},
MastNode::External(node) => {
self.eq_hash_from_parts(node.before_enter(), node.after_exit(), &[], node.digest())
},
}
}

fn eq_hash_from_parts(
&self,
before_enter_ids: &[DecoratorId],
after_exit_ids: &[DecoratorId],
children_ids: &[MastNodeId],
node_digest: RpoDigest,
) -> EqHash {
let pre_decorator_hash_bytes =
before_enter_ids.iter().flat_map(|&id| self[id].eq_hash().as_bytes());
let post_decorator_hash_bytes =
after_exit_ids.iter().flat_map(|&id| self[id].eq_hash().as_bytes());

// Reminder: the `EqHash`'s decorator root will be `None` if and only if there are no
// decorators attached to the node, and all children have no decorator roots (meaning that
// there are no decorators in all the descendants).
if pre_decorator_hash_bytes.clone().next().is_none()
&& post_decorator_hash_bytes.clone().next().is_none()
&& children_ids
.iter()
.filter_map(|child_id| self.hash_by_node_id[child_id].decorator_root)
.next()
.is_none()
{
EqHash::new(node_digest)
} else {
let children_decorator_roots = children_ids
.iter()
.filter_map(|child_id| self.hash_by_node_id[child_id].decorator_root)
.flat_map(|decorator_root| decorator_root.as_bytes());
let decorator_bytes_to_hash: Vec<u8> = pre_decorator_hash_bytes
.chain(post_decorator_hash_bytes)
.chain(children_decorator_roots)
.collect();

let decorator_root = Blake3_256::hash(&decorator_bytes_to_hash);
EqHash::with_decorator_root(node_digest, decorator_root)
}
EqHash::from_mast_node(&self.mast_forest, &self.hash_by_node_id, node)
}
}

Expand Down Expand Up @@ -582,33 +476,6 @@ impl IndexMut<DecoratorId> for MastForestBuilder {
}
}

// EQ HASH
// ================================================================================================

/// Represents the hash used to test for equality between [`MastNode`]s.
///
/// The decorator root will be `None` if and only if there are no decorators attached to the node,
/// and all children have no decorator roots (meaning that there are no decorators in all the
/// descendants).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct EqHash {
mast_root: RpoDigest,
decorator_root: Option<Blake3Digest<32>>,
}

impl EqHash {
fn new(mast_root: RpoDigest) -> Self {
Self { mast_root, decorator_root: None }
}

fn with_decorator_root(mast_root: RpoDigest, decorator_root: Blake3Digest<32>) -> Self {
Self {
mast_root,
decorator_root: Some(decorator_root),
}
}
}

// HELPER FUNCTIONS
// ================================================================================================

Expand Down
73 changes: 73 additions & 0 deletions assembly/src/assembler/mast_forest_merger_tests.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
use miette::{IntoDiagnostic, Report};
use vm_core::mast::{MastForest, MastForestRootMap};

use crate::{testing::TestContext, Assembler};

#[allow(clippy::type_complexity)]
fn merge_programs(
program_a: &str,
program_b: &str,
) -> Result<(MastForest, MastForest, MastForest, MastForestRootMap), Report> {
let context = TestContext::new();
let module = context.parse_module_with_path("lib::mod".parse().unwrap(), program_a)?;

let lib_a = Assembler::new(context.source_manager()).assemble_library([module])?;

let mut assembler = Assembler::new(context.source_manager());
assembler.add_library(lib_a.clone())?;
let lib_b = assembler.assemble_library([program_b])?.mast_forest().as_ref().clone();
let lib_a = lib_a.mast_forest().as_ref().clone();

let (merged, root_maps) = MastForest::merge([&lib_a, &lib_b]).into_diagnostic()?;

Ok((lib_a, lib_b, merged, root_maps))
}

/// Tests that an assembler-produced library's forests can be merged and that external nodes are
/// replaced by their referenced procedures.
#[test]
fn mast_forest_merge_assembler() {
let lib_a = r#"
export.foo
push.19
end

export.qux
swap drop
end
"#;

let lib_b = r#"
use.lib::mod

export.qux_duplicate
swap drop
end

export.bar
push.2
if.true
push.3
else
while.true
add
push.23
end
end
exec.mod::foo
end"#;

let (forest_a, forest_b, merged, root_maps) = merge_programs(lib_a, lib_b).unwrap();

for (forest_idx, forest) in [forest_a, forest_b].into_iter().enumerate() {
for root in forest.procedure_roots() {
let original_digest = forest.nodes()[root.as_usize()].digest();
let new_root = root_maps.map_root(forest_idx, root).unwrap();
let new_digest = merged.nodes()[new_root.as_usize()].digest();
assert_eq!(original_digest, new_digest);
}
}

// Assert that the external node for the import was removed during merging.
merged.nodes().iter().for_each(|node| assert!(!node.is_external()));
}
3 changes: 3 additions & 0 deletions assembly/src/assembler/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,9 @@ mod procedure;
#[cfg(test)]
mod tests;

#[cfg(test)]
mod mast_forest_merger_tests;

use self::{
basic_block_builder::BasicBlockBuilder,
module_graph::{CallerInfo, ModuleGraph, ResolvedTarget},
Expand Down
Loading
Loading