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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
/*
 * AquaVM Workflow Engine
 *
 * Copyright (C) 2024 Fluence DAO
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation version 3 of the
 * License.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

pub(crate) mod errors;
pub(crate) mod stream_map_key;

use self::stream_map_key::StreamMapKey;
use crate::execution_step::value_types::StreamMap;
use crate::execution_step::ExecutionResult;
use crate::execution_step::Generation;
use crate::execution_step::ValueAggregate;

use air_parser::ast::Span;
use air_parser::AirPos;
use air_trace_handler::TraceHandler;

use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::collections::HashMap;
use std::fmt;

// TODO This module should be unified with its Stream counterpart.
pub(crate) struct StreamMapValueDescriptor<'stream_name> {
    pub value: ValueAggregate,
    pub name: &'stream_name str,
    pub generation: Generation,
    pub position: AirPos,
}

impl<'stream_name> StreamMapValueDescriptor<'stream_name> {
    pub fn new(value: ValueAggregate, name: &'stream_name str, generation: Generation, position: AirPos) -> Self {
        Self {
            value,
            name,
            generation,
            position,
        }
    }
}

pub(crate) struct StreamMapDescriptor {
    pub(super) span: Span,
    pub(super) stream_map: StreamMap,
}

impl StreamMapDescriptor {
    pub(super) fn global(stream_map: StreamMap) -> Self {
        Self {
            span: Span::new(0.into(), usize::MAX.into()),
            stream_map,
        }
    }

    pub(super) fn restricted(stream_map: StreamMap, span: Span) -> Self {
        Self { span, stream_map }
    }
}

impl fmt::Display for StreamMapDescriptor {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, " <{}> - <{}>: {}", self.span.left, self.span.right, self.stream_map)
    }
}

pub(super) fn find_closest<'d>(
    descriptors: impl DoubleEndedIterator<Item = &'d StreamMapDescriptor>,
    position: AirPos,
) -> Option<&'d StreamMap> {
    // descriptors are placed in a order of decreasing scopes, so it's enough to get the latest suitable
    descriptors
        .rev()
        .find(|d| d.span.contains_position(position))
        .map(|d| &d.stream_map)
}

pub(super) fn find_closest_mut<'d>(
    descriptors: impl DoubleEndedIterator<Item = &'d mut StreamMapDescriptor>,
    position: AirPos,
) -> Option<&'d mut StreamMap> {
    // descriptors are placed in a order of decreasing scopes, so it's enough to get the latest suitable
    descriptors
        .rev()
        .find(|d| d.span.contains_position(position))
        .map(|d| &mut d.stream_map)
}
#[derive(Default)]
pub(crate) struct StreamMaps {
    // this one is optimized for speed (not for memory), because it's unexpected
    // that a script could have a lot of new.
    // TODO: use shared string (Rc<String>) to avoid copying.
    stream_maps: HashMap<String, Vec<StreamMapDescriptor>>,
}

impl StreamMaps {
    pub(crate) fn get(&self, name: &str, position: AirPos) -> Option<&StreamMap> {
        self.stream_maps
            .get(name)
            .and_then(|descriptors| find_closest(descriptors.iter(), position))
    }

    pub(crate) fn get_mut(&mut self, name: &str, position: AirPos) -> Option<&mut StreamMap> {
        self.stream_maps
            .get_mut(name)
            .and_then(|descriptors| find_closest_mut(descriptors.iter_mut(), position))
    }

    pub(crate) fn add_stream_map_value(
        &mut self,
        key: StreamMapKey,
        value_descriptor: StreamMapValueDescriptor<'_>,
    ) -> ExecutionResult<()> {
        let StreamMapValueDescriptor {
            value,
            name,
            generation,
            position,
        } = value_descriptor;

        match self.get_mut(name, position) {
            Some(stream_map) => stream_map.insert(key, &value, generation),
            None => {
                // streams could be created in three ways:
                //  - after met new instruction with stream name that isn't present in streams
                //    (it's the only way to create restricted streams)
                //  - by calling add_global_stream with generation that come from data
                //    for global streams
                //  - and by this function, and if there is no such a streams in streams,
                //    it means that a new global one should be created.
                let mut stream_map = StreamMap::new();
                stream_map.insert(key, &value, generation)?;
                let descriptor = StreamMapDescriptor::global(stream_map);
                self.stream_maps.insert(name.to_string(), vec![descriptor]);
                Ok(())
            }
        }
    }

    pub(crate) fn meet_scope_start(&mut self, name: impl Into<String>, span: Span) {
        let name = name.into();

        let new_stream_map = StreamMap::new();
        let new_descriptor = StreamMapDescriptor::restricted(new_stream_map, span);
        match self.stream_maps.entry(name) {
            Occupied(mut entry) => {
                entry.get_mut().push(new_descriptor);
            }
            Vacant(entry) => {
                entry.insert(vec![new_descriptor]);
            }
        }
    }

    pub(crate) fn meet_scope_end(&mut self, name: String, trace_ctx: &mut TraceHandler) -> ExecutionResult<()> {
        // unwraps are safe here because met_scope_end must be called after met_scope_start
        let stream_map_descriptors = self.stream_maps.get_mut(&name).unwrap();
        // delete a stream after exit from a scope
        let mut last_descriptor = stream_map_descriptors.pop().unwrap();
        if stream_map_descriptors.is_empty() {
            // streams should contain only non-empty stream embodiments
            self.stream_maps.remove(&name);
        }

        last_descriptor.stream_map.compactify(trace_ctx)
    }

    pub(crate) fn compactify(&mut self, trace_ctx: &mut TraceHandler) -> ExecutionResult<()> {
        for (_, descriptors) in self.stream_maps.iter_mut() {
            for descriptor in descriptors.iter_mut() {
                descriptor.stream_map.compactify(trace_ctx)?;
            }
        }

        Ok(())
    }
}

impl fmt::Display for StreamMaps {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (name, descriptors) in self.stream_maps.iter() {
            if let Some(last_descriptor) = descriptors.last() {
                writeln!(f, "{name} => {last_descriptor}")?;
            }
        }
        Ok(())
    }
}