sway_ir/
storage_key.rs

1//! A value representing a storage key. Every storage field access in the program
2//! corresponds to a [StorageKey].
3
4use std::vec;
5
6use crate::{
7    context::Context, irtype::Type, pretty::DebugWithContext, Constant, ConstantContent,
8    ConstantValue, B256,
9};
10
11/// A wrapper around an [ECS](https://github.com/orlp/slotmap) handle into the
12/// [`Context`].
13#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, DebugWithContext)]
14pub struct StorageKey(#[in_context(storage_keys)] pub slotmap::DefaultKey);
15
16#[doc(hidden)]
17#[derive(Clone, DebugWithContext)]
18pub struct StorageKeyContent {
19    pub ptr_ty: Type,
20    pub key: Constant,
21}
22
23impl StorageKey {
24    pub fn new(context: &mut Context, slot: [u8; 32], offset: u64, field_id: [u8; 32]) -> Self {
25        // Construct `ptr { b256, u64, b256 }`.
26        let b256_ty = Type::get_b256(context);
27        let uint64_ty = Type::get_uint64(context);
28
29        let key_ty = Type::new_struct(context, vec![b256_ty, uint64_ty, b256_ty]);
30        let ptr_ty = Type::new_typed_pointer(context, key_ty);
31
32        let slot = ConstantContent::new_b256(context, slot);
33        let offset = ConstantContent::new_uint(context, 64, offset);
34        let field_id = ConstantContent::new_b256(context, field_id);
35
36        let key = ConstantContent::new_struct(
37            context,
38            vec![b256_ty, uint64_ty, b256_ty],
39            vec![slot, offset, field_id],
40        );
41
42        let key = Constant::unique(context, key);
43
44        let content = StorageKeyContent { ptr_ty, key };
45
46        StorageKey(context.storage_keys.insert(content))
47    }
48
49    /// Return the storage key type, which is always `ptr { b256, u64, b256 }`.
50    pub fn get_type(&self, context: &Context) -> Type {
51        context.storage_keys[self.0].ptr_ty
52    }
53
54    /// Return the storage key, which is a constant of type `{ b256, u64, b256 }`.
55    pub fn get_key(&self, context: &Context) -> Constant {
56        context.storage_keys[self.0].key
57    }
58
59    /// Return the three parts of this storage key: `(slot, offset, field_id)`.
60    pub fn get_parts<'a>(&self, context: &'a Context) -> (&'a B256, u64, &'a B256) {
61        let ConstantContent {
62            value: ConstantValue::Struct(fields),
63            ..
64        } = &context.storage_keys[self.0].key.get_content(context)
65        else {
66            unreachable!("`StorageKey::key` constant content is a struct with three fields");
67        };
68
69        let ConstantContent {
70            value: ConstantValue::B256(slot),
71            ..
72        } = &fields[0]
73        else {
74            unreachable!("storage key slot is a `B256` constant");
75        };
76        let ConstantContent {
77            value: ConstantValue::Uint(offset),
78            ..
79        } = &fields[1]
80        else {
81            unreachable!("storage key offset is a `u64` constant");
82        };
83        let ConstantContent {
84            value: ConstantValue::B256(field_id),
85            ..
86        } = &fields[2]
87        else {
88            unreachable!("storage key field_id is a `B256` constant");
89        };
90
91        (slot, *offset, field_id)
92    }
93}