weval 0.4.1

The WebAssembly partial evaluator
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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

//! Caching of weval results.

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::path::Path;

pub type ModuleHash = [u8; 32]; // SHA-256 hash.

pub(crate) fn compute_hash(raw_bytes: &[u8]) -> ModuleHash {
    Sha256::digest(raw_bytes).into()
}

/// Cache result: compiled Wasm bytecode, with signature.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct CacheData {
    /// Function signature index.
    pub sig: u32,
    /// Function name.
    pub name: String,
    /// Raw function body bytecode, including locals.
    pub body: Vec<u8>,
}

pub(crate) struct Cache {
    module_hash: ModuleHash,
    db: Option<sqlite::ConnectionThreadSafe>,
    db_ro: Option<sqlite::ConnectionThreadSafe>,
}

pub(crate) struct CacheThreadCtx<'a> {
    cache: &'a Cache,
    lookup_stmt: Option<sqlite::Statement<'a>>,
    insert_stmt: Option<sqlite::Statement<'a>>,
    ro_lookup_stmt: Option<sqlite::Statement<'a>>,
}

impl Cache {
    pub fn open(
        path: Option<&Path>,
        path_ro: Option<&Path>,
        module_hash: ModuleHash,
    ) -> anyhow::Result<Cache> {
        let db = match path {
            Some(path) => {
                let db = sqlite::Connection::open_thread_safe(path)?;
                db.execute(
                    r#"
                    CREATE TABLE IF NOT EXISTS weval_cache(
                        module_hash BLOB NOT NULL,
                        key BLOB NOT NULL,
                        result BLOB NOT NULL,
                        created_time INTEGER NOT NULL
                     );
                     CREATE INDEX IF NOT EXISTS idx ON weval_cache(
                         module_hash, key
                     );
                     "#,
                )?;
                Some(db)
            }
            None => None,
        };
        let db_ro = match path_ro {
            Some(path_ro) => Some(sqlite::Connection::open_thread_safe_with_flags(
                path_ro,
                sqlite::OpenFlags::new().with_read_only(),
            )?),
            None => None,
        };
        Ok(Cache {
            module_hash,
            db,
            db_ro,
        })
    }

    pub fn can_insert(&self) -> bool {
        self.db.is_some()
    }

    pub fn thread(&self) -> anyhow::Result<CacheThreadCtx<'_>> {
        let lookup_stmt = match self.db.as_ref() {
            Some(db) => Some(db.prepare(
                r#"
                SELECT result FROM weval_cache
                    WHERE module_hash=? AND key=?
                "#,
            )?),
            None => None,
        };
        let ro_lookup_stmt = match self.db_ro.as_ref() {
            Some(db_ro) => Some(db_ro.prepare(
                r#"
                SELECT result FROM weval_cache
                    WHERE module_hash=? AND key=?
                "#,
            )?),
            None => None,
        };
        let insert_stmt = match self.db.as_ref() {
            Some(db) => Some(db.prepare(
                r#"
                INSERT INTO weval_cache
                    (module_hash, key, result, created_time)
                VALUES
                    (?, ?, ?, unixepoch())
                "#,
            )?),
            None => None,
        };
        Ok(CacheThreadCtx {
            cache: self,
            lookup_stmt,
            insert_stmt,
            ro_lookup_stmt,
        })
    }
}

impl<'a> CacheThreadCtx<'a> {
    pub fn lookup(&mut self, key: &[u8]) -> anyhow::Result<Option<CacheData>> {
        let mut result = None;
        for lookup in self
            .ro_lookup_stmt
            .iter_mut()
            .chain(self.lookup_stmt.iter_mut())
        {
            lookup.bind((1, &self.cache.module_hash[..]))?;
            lookup.bind((2, key))?;

            while lookup.next()? == sqlite::State::Row {
                let data: Vec<u8> = lookup.read(0)?;
                result = Some(bincode::deserialize(&data)?);
            }

            lookup.reset()?;
            if result.is_some() {
                break;
            }
        }
        Ok(result)
    }

    pub fn insert(&mut self, key: &[u8], data: CacheData) -> anyhow::Result<()> {
        if let Some(insert) = self.insert_stmt.as_mut() {
            let data = bincode::serialize(&data)?;
            insert.bind((1, &self.cache.module_hash[..]))?;
            insert.bind((2, key))?;
            insert.bind((3, &data[..]))?;
            while insert.next()? == sqlite::State::Row {}
            insert.reset()?;
        }
        Ok(())
    }
}