1use std::vec;
5
6use crate::{
7 context::Context, irtype::Type, pretty::DebugWithContext, Constant, ConstantContent,
8 ConstantValue, B256,
9};
10
11#[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 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 pub fn get_type(&self, context: &Context) -> Type {
51 context.storage_keys[self.0].ptr_ty
52 }
53
54 pub fn get_key(&self, context: &Context) -> Constant {
56 context.storage_keys[self.0].key
57 }
58
59 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}