[framework-snaps] Update snapshot loading

crates/sui-framework-snapshot/src/lib.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use serde::{Deserialize, Serialize};
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use std::{fs, io::Read, path::PathBuf};
 use sui_framework::{SystemPackage, SystemPackageMetadata};
 use sui_types::base_types::ObjectID;
@@ -100,8 +100,7 @@ pub fn update_bytecode_snapshot_manifest(
 }
 
 pub fn load_bytecode_snapshot(protocol_version: u64) -> anyhow::Result<Vec<SystemPackage>> {
-    let mut snapshot_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
-    snapshot_path.extend(["bytecode_snapshot", protocol_version.to_string().as_str()]);
+    let snapshot_path = snapshot_path_for_version(protocol_version)?;
     let mut snapshots: BTreeMap<ObjectID, SystemPackage> = fs::read_dir(&snapshot_path)?
         .flatten()
         .map(|entry| {
@@ -128,3 +127,33 @@ pub fn load_bytecode_snapshot(protocol_version: u64) -> anyhow::Result<Vec<Syste
 pub fn manifest_path() -> PathBuf {
     PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("manifest.json")
 }
+
+/// Given a protocol version:
+/// * The path to the snapshot directory for that version is returned, if it exists.
+/// * If the version is greater than the latest snapshot version, then `Ok(None)` is returned.
+/// * If the version does not exist, but there are snapshots present with versions greater than
+///   `version`, then the smallest snapshot number greater than `version` is returned.
+fn snapshot_path_for_version(version: u64) -> anyhow::Result<PathBuf> {
+    let snapshot_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("bytecode_snapshot");
+    let mut snapshots = BTreeSet::new();
+
+    for entry in fs::read_dir(&snapshot_dir)? {
+        let entry = entry?;
+        let path = entry.path();
+        if path.is_dir() {
+            if let Some(snapshot_number) = path
+                .file_name()
+                .and_then(|n| n.to_str())
+                .and_then(|n| n.parse::<u64>().ok())
+            {
+                snapshots.insert(snapshot_number);
+            }
+        }
+    }
+
+    snapshots
+        .range(version..)
+        .next()
+        .map(|v| snapshot_dir.join(v.to_string()))
+        .ok_or_else(|| anyhow::anyhow!("No snapshot found for version {}", version))
+}