1
use std::path::PathBuf;
2

            
3
use anyhow::Result;
4
use bytes::Bytes;
5
use risc0_binfmt::{MemoryImage, Program};
6
use risc0_zkvm::{GUEST_MAX_MEM, PAGE_SIZE};
7
use tokio::fs::read;
8

            
9
pub struct Image {
10
    pub id: String,
11
    pub data: Option<Program>,
12
    bytes: Option<Bytes>,
13
    pub size: u64,
14
    pub path: PathBuf,
15
    pub last_used: u64,
16
}
17

            
18
impl Image {
19
    fn load_elf(elf: &[u8]) -> Result<Program> {
20
        let program = Program::load_elf(elf, GUEST_MAX_MEM as u32)?;
21
        Ok(program)
22
    }
23

            
24
    fn mem_img(program: &Program) -> Result<MemoryImage> {
25
        let image = MemoryImage::new(program, PAGE_SIZE as u32)?;
26
        Ok(image)
27
    }
28

            
29
    pub fn bytes(&self) -> Option<&Bytes> {
30
        self.bytes.as_ref()
31
    }
32

            
33
    pub fn from_bytes(bytes: Bytes) -> Result<Image> {
34
        let program = Image::load_elf(&bytes)?;
35
        let img = Image::mem_img(&program)?;
36
        Ok(Image {
37
            id: img.compute_id().to_string(),
38
            bytes: Some(bytes),
39
            data: Some(program),
40
            size: img.pages.len() as u64 * PAGE_SIZE as u64,
41
            path: PathBuf::new(),
42
            last_used: 0,
43
        })
44
    }
45

            
46
    pub async fn new(path: PathBuf) -> Result<Image> {
47
        let data = read(&path).await?;
48
        let program = Image::load_elf(&data)?;
49
        let img = Image::mem_img(&program)?;
50

            
51
        Ok(Image {
52
            id: img.compute_id().to_string(),
53
            bytes: Some(Bytes::from(data)),
54
            data: Some(program),
55
            size: img.pages.len() as u64 * PAGE_SIZE as u64,
56
            path,
57
            last_used: 0,
58
        })
59
    }
60

            
61
    pub fn compress(&mut self) {
62
        self.data = None;
63
        self.bytes = None;
64
    }
65

            
66
    pub async fn load(&mut self) -> Result<()> {
67
        if self.data.is_some() {
68
            return Ok(());
69
        }
70
        let data = read(&self.path).await?;
71
        let program = Image::load_elf(&data)?;
72
        self.data = Some(program);
73
        self.bytes = Some(Bytes::from(data));
74
        Ok(())
75
    }
76

            
77
    pub fn get_memory_image(&self) -> Result<MemoryImage> {
78
        let program = self
79
            .data
80
            .as_ref()
81
            .ok_or_else(|| anyhow::anyhow!("No data"))?;
82
        let image = Image::mem_img(program)?;
83
        Ok(image)
84
    }
85
}