1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

// Copyright (c) The Diem Core Contributors
// Copyright (c) The Move Contributors
// SPDX-License-Identifier: Apache-2.0

use anyhow::{anyhow, Result};
use move_binary_format::CompiledModule;
use move_core_types::{language_storage::ModuleId, resolver::ModuleResolver};
use std::{
    borrow::Borrow,
    cell::RefCell,
    collections::{btree_map::Entry, BTreeMap},
    fmt::Debug,
    sync::{Arc, RwLock},
};

/// A persistent storage that can fetch the bytecode for a given module id
/// TODO: do we want to implement this in a way that allows clients to cache struct layouts?
pub trait GetModule {
    type Error: Debug;
    type Item: Borrow<CompiledModule>;

    fn get_module_by_id(&self, id: &ModuleId) -> Result<Option<Self::Item>, Self::Error>;
}

/// Simple in-memory module cache
pub struct ModuleCache<R: ModuleResolver> {
    cache: RefCell<BTreeMap<ModuleId, CompiledModule>>,
    resolver: R,
}

impl<R: ModuleResolver> ModuleCache<R> {
    pub fn new(resolver: R) -> Self {
        ModuleCache {
            cache: RefCell::new(BTreeMap::new()),
            resolver,
        }
    }

    pub fn add(&self, id: ModuleId, m: CompiledModule) {
        self.cache.borrow_mut().insert(id, m);
    }
}

impl<R: ModuleResolver> GetModule for ModuleCache<R> {
    type Error = anyhow::Error;
    type Item = CompiledModule;

    fn get_module_by_id(&self, id: &ModuleId) -> Result<Option<CompiledModule>, Self::Error> {
        Ok(Some(match self.cache.borrow_mut().entry(id.clone()) {
            Entry::Vacant(entry) => {
                let module_bytes = self
                    .resolver
                    .get_module(id)
                    .map_err(|_| anyhow!("Failed to get module {:?}", id))?
                    .ok_or_else(|| anyhow!("Module {:?} doesn't exist", id))?;
                let module = CompiledModule::deserialize(&module_bytes)
                    .map_err(|_| anyhow!("Failure deserializing module {:?}", id))?;
                entry.insert(module.clone());
                module
            }
            Entry::Occupied(entry) => entry.get().clone(),
        }))
    }
}

/// Simple in-memory module cache that implements Sync
pub struct SyncModuleCache<R: ModuleResolver> {
    cache: RwLock<BTreeMap<ModuleId, Arc<CompiledModule>>>,
    resolver: R,
}

impl<R: ModuleResolver> SyncModuleCache<R> {
    pub fn new(resolver: R) -> Self {
        SyncModuleCache {
            cache: RwLock::new(BTreeMap::new()),
            resolver,
        }
    }

    pub fn add(&self, id: ModuleId, m: CompiledModule) {
        self.cache.write().unwrap().insert(id, Arc::new(m));
    }
}

impl<R: ModuleResolver> GetModule for SyncModuleCache<R> {
    type Error = anyhow::Error;
    type Item = Arc<CompiledModule>;

    fn get_module_by_id(&self, id: &ModuleId) -> Result<Option<Arc<CompiledModule>>, Self::Error> {
        if let Some(compiled_module) = self.cache.read().unwrap().get(id) {
            return Ok(Some(compiled_module.clone()));
        }

        if let Some(module_bytes) = self
            .resolver
            .get_module(id)
            .map_err(|_| anyhow!("Failed to get module {:?}", id))?
        {
            let module = Arc::new(
                CompiledModule::deserialize(&module_bytes)
                    .map_err(|_| anyhow!("Failure deserializing module {:?}", id))?,
            );

            self.cache
                .write()
                .unwrap()
                .insert(id.clone(), module.clone());
            Ok(Some(module))
        } else {
            Ok(None)
        }
    }
}