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

//! SubIterator for a Component

use crate::ir::id::{FunctionID, ModuleID};
use crate::ir::types::Location;
use crate::subiterator::module_subiterator::ModuleSubIterator;
use std::collections::HashMap;

/// Sub-iterator for a Component. Keeps track of current location in a Component.
pub struct ComponentSubIterator {
    /// The current module the SubIterator is at.
    curr_mod: ModuleID,
    /// The total number of modules in the component.
    num_mods: usize,
    /// The module iterator used to keep track of the location in the module.
    pub(crate) mod_iterator: ModuleSubIterator,
    /// Metadata that maps Module Index -> Function Index -> Instruction Index
    metadata: HashMap<ModuleID, HashMap<FunctionID, usize>>,
    /// Map of Module -> Functions to skip in that module. Provide an empty HashMap if no functions are to be skipped.
    skip_funcs: HashMap<ModuleID, Vec<FunctionID>>,
}

impl ComponentSubIterator {
    /// Creates a new ComponentSubIterator
    pub fn new(
        curr_mod: ModuleID,
        num_mods: usize,
        metadata: HashMap<ModuleID, HashMap<FunctionID, usize>>,
        skip_funcs: HashMap<ModuleID, Vec<FunctionID>>,
    ) -> Self {
        // Get current skip func
        // initializes to the first module
        ComponentSubIterator {
            curr_mod,
            num_mods,
            metadata: metadata.clone(),
            mod_iterator: ModuleSubIterator::new(
                metadata.get(&ModuleID(0)).unwrap().keys().len() as u32,
                (*metadata.get(&curr_mod).unwrap()).clone(),
                match skip_funcs.contains_key(&curr_mod) {
                    true => skip_funcs.get(&curr_mod).unwrap().clone(),
                    false => vec![],
                },
            ),
            skip_funcs,
        }
    }

    /// Resets the ComponentSubIterator and all child SubIterators
    pub fn reset(&mut self) {
        *self.curr_mod = 0;
        self.mod_iterator
            .reset_from_comp_iterator((*self.metadata.get(&self.curr_mod).unwrap()).clone());
    }

    /// Goes to the next module enclosed by the component
    fn next_module(&mut self) -> bool {
        *self.curr_mod += 1;
        if *self.curr_mod < self.num_mods as u32 {
            let num_funcs = self.metadata.get(&self.curr_mod).unwrap().keys().len() as u32;
            let met = self.metadata.get(&self.curr_mod).unwrap().clone();
            // If we're defining a new module, we have to reset function
            self.mod_iterator = ModuleSubIterator::new(
                num_funcs,
                met,
                match self.skip_funcs.contains_key(&self.curr_mod) {
                    true => self.skip_funcs.get(&self.curr_mod).unwrap().clone(),
                    false => vec![],
                },
            );
            true
        } else {
            false
        }
    }

    /// Gets the index of the current module in the component
    pub fn curr_mod_idx(&self) -> ModuleID {
        self.curr_mod
    }

    /// Gets the index of the current function in the current module
    pub fn curr_func_idx(&self) -> FunctionID {
        self.mod_iterator.curr_func
    }

    /// Gets the index of the current instruction in the current function
    pub fn curr_instr_idx(&self) -> usize {
        self.mod_iterator.func_iterator.curr_instr
    }

    /// Checks if the SubIterator has finished traversing all the modules
    pub fn end(&self) -> bool {
        *self.curr_mod as usize == self.num_mods
    }

    /// Returns the Current Location as a Location and a bool value that
    /// says whether the location is at the end of the function.
    pub fn curr_loc(&self) -> (Location, bool) {
        if let (
            Location::Module {
                func_idx,
                instr_idx,
            },
            is_end,
        ) = self.mod_iterator.curr_loc()
        {
            (
                Location::Component {
                    mod_idx: self.curr_mod,
                    func_idx,
                    instr_idx,
                },
                is_end,
            )
        } else {
            panic!("Should have gotten Module Location from Module Iterator and not Component Location!");
        }
    }

    #[allow(clippy::should_implement_trait)]
    /// Goes to the next instruction in the component
    pub fn next(&mut self) -> bool {
        if self.mod_iterator.has_next() {
            self.mod_iterator.next()
        } else {
            self.next_module()
        }
    }
}