sway_ir/
module.rs

1//! A scope containing a collection of [`Function`]s and constant values.
2//!
3//! A module also has a 'kind' corresponding to the different Sway module types.
4
5use std::{cell::Cell, collections::BTreeMap};
6
7use crate::{
8    context::Context,
9    function::{Function, FunctionIterator},
10    Constant, GlobalVar, MetadataIndex, Type,
11};
12
13/// A wrapper around an [ECS](https://github.com/orlp/slotmap) handle into the
14/// [`Context`].
15#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
16pub struct Module(pub slotmap::DefaultKey);
17
18#[doc(hidden)]
19pub struct ModuleContent {
20    pub kind: Kind,
21    pub functions: Vec<Function>,
22    pub global_variables: BTreeMap<Vec<String>, GlobalVar>,
23    pub configs: BTreeMap<String, ConfigContent>,
24}
25
26#[derive(Clone, Debug)]
27pub enum ConfigContent {
28    V0 {
29        name: String,
30        ty: Type,
31        ptr_ty: Type,
32        constant: Constant,
33        opt_metadata: Option<MetadataIndex>,
34    },
35    V1 {
36        name: String,
37        ty: Type,
38        ptr_ty: Type,
39        encoded_bytes: Vec<u8>,
40        decode_fn: Cell<Function>,
41        opt_metadata: Option<MetadataIndex>,
42    },
43}
44
45/// The different 'kinds' of Sway module: `Contract`, `Library`, `Predicate` or `Script`.
46#[derive(Clone, Copy, Debug, Eq, PartialEq)]
47pub enum Kind {
48    Contract,
49    Library,
50    Predicate,
51    Script,
52}
53
54impl Module {
55    /// Return a new module of a specific kind.
56    pub fn new(context: &mut Context, kind: Kind) -> Module {
57        let content = ModuleContent {
58            kind,
59            functions: Vec::new(),
60            global_variables: BTreeMap::new(),
61            configs: BTreeMap::new(),
62        };
63        Module(context.modules.insert(content))
64    }
65
66    /// Get this module's [`Kind`].
67    pub fn get_kind(&self, context: &Context) -> Kind {
68        context.modules[self.0].kind
69    }
70
71    /// Return an iterator over each of the [`Function`]s in this module.
72    pub fn function_iter(&self, context: &Context) -> FunctionIterator {
73        FunctionIterator::new(context, self)
74    }
75
76    /// Add a global variable value to this module.
77    pub fn add_global_variable(
78        &self,
79        context: &mut Context,
80        call_path: Vec<String>,
81        const_val: GlobalVar,
82    ) {
83        context.modules[self.0]
84            .global_variables
85            .insert(call_path, const_val);
86    }
87
88    /// Get a named global variable from this module, if found.
89    pub fn get_global_variable(
90        &self,
91        context: &Context,
92        call_path: &Vec<String>,
93    ) -> Option<GlobalVar> {
94        context.modules[self.0]
95            .global_variables
96            .get(call_path)
97            .copied()
98    }
99
100    /// Lookup global variable name
101    pub fn lookup_global_variable_name(
102        &self,
103        context: &Context,
104        global: &GlobalVar,
105    ) -> Option<String> {
106        context.modules[self.0]
107            .global_variables
108            .iter()
109            .find(|(_key, val)| *val == global)
110            .map(|(key, _)| key.join("::"))
111    }
112
113    /// Add a config value to this module.
114    pub fn add_config(&self, context: &mut Context, name: String, content: ConfigContent) {
115        context.modules[self.0].configs.insert(name, content);
116    }
117
118    /// Get a named config content from this module, if found.
119    pub fn get_config<'a>(&self, context: &'a Context, name: &str) -> Option<&'a ConfigContent> {
120        context.modules[self.0].configs.get(name)
121    }
122
123    /// Removed a function from the module.  Returns true if function was found and removed.
124    ///
125    /// **Use with care!  Be sure the function is not an entry point nor called at any stage.**
126    pub fn remove_function(&self, context: &mut Context, function: &Function) {
127        context
128            .modules
129            .get_mut(self.0)
130            .expect("Module must exist in context.")
131            .functions
132            .retain(|mod_fn| mod_fn != function);
133    }
134
135    pub fn iter_configs<'a>(
136        &'a self,
137        context: &'a Context,
138    ) -> impl Iterator<Item = &'a ConfigContent> + 'a {
139        context.modules[self.0].configs.values()
140    }
141}
142
143/// An iterator over [`Module`]s within a [`Context`].
144pub struct ModuleIterator {
145    modules: Vec<slotmap::DefaultKey>,
146    next: usize,
147}
148
149impl ModuleIterator {
150    /// Return a new [`Module`] iterator.
151    pub fn new(context: &Context) -> ModuleIterator {
152        // Copy all the current modules indices, so they may be modified in the context during
153        // iteration.
154        ModuleIterator {
155            modules: context.modules.iter().map(|pair| pair.0).collect(),
156            next: 0,
157        }
158    }
159}
160
161impl Iterator for ModuleIterator {
162    type Item = Module;
163
164    fn next(&mut self) -> Option<Module> {
165        if self.next < self.modules.len() {
166            let idx = self.next;
167            self.next += 1;
168            Some(Module(self.modules[idx]))
169        } else {
170            None
171        }
172    }
173}