1  use std::{
  2     cell::RefCell,
  3     collections::VecDeque,
  4     fmt::{
  5        self,
  6        Write as _,
  7     },
  8     ops::{
  9        self,
 10        Add as _,
 11     },
 12  };
 13  
 14  use derive_more::{
 15     Deref,
 16     DerefMut,
 17  };
 18  use dup::Dupe as _;
 19  use ranged::Span;
 20  use ust::{
 21     COLORS,
 22     Display,
 23     INDENT_WIDTH,
 24     STYLE_GUTTER,
 25     Write,
 26     style::{
 27        self,
 28        StyledExt as _,
 29     },
 30     terminal::{
 31        self,
 32        DOT,
 33        LEFT_TO_RIGHT,
 34        RIGHT_TO_BOTTOM,
 35        TOP_TO_BOTTOM,
 36     },
 37     with,
 38     write,
 39  };
 40  
 41  use crate::{
 42     Argument,
 43     Operation,
 44     Value,
 45     value,
 46  };
 47  
 48  const ENCODED_U64_LEN_MAX: usize = 9;
 49  const ENCODED_U16_LEN_MAX: usize = 0_u16.to_le_bytes().len();
 50  const ENCODED_OPERATION_LEN: usize = 1;
 51  
 52  #[derive(Deref, DerefMut, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
 53  pub struct ByteIndex(usize);
 54  
 55  impl ByteIndex {
 56     #[must_use]
 57     pub fn dummy() -> Self {
 58        Self(usize::MAX)
 59     }
 60  }
 61  
 62  #[derive(Deref, DerefMut, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
 63  pub struct ValueIndex(usize);
 64  
 65  impl ValueIndex {
 66     #[must_use]
 67     pub fn dummy() -> Self {
 68        Self(usize::MAX)
 69     }
 70  }
 71  
 72  pub struct Code {
 73     bytes: Vec<u8>,
 74  
 75     path:  value::Path,
 76     spans: Vec<(ByteIndex, Span)>,
 77  
 78     values: Vec<Value>,
 79  }
 80  
 81  impl Display for Code {
 82     fn display_styled(&self, writer: &mut dyn Write) -> fmt::Result {
 83        const STYLE_JUMP_ADDRESS: style::Style = style::Color::BrightYellow.fg().bold().underline();
 84  
 85        enum CodeType {
 86           Suspend,
 87           Lambda,
 88        }
 89  
 90        let mut codes = VecDeque::from([(0_usize, CodeType::Suspend, self)]);
 91  
 92        while let Some((code_index, code_type, code)) = codes.pop_back() {
 93           let highlighted = RefCell::new(Vec::<ByteIndex>::new());
 94  
 95           // INDENT: "0x123 | "
 96           let index_width = 2 + terminal::number_hex_width(code.bytes.len() - 1);
 97           let indent_index = RefCell::new(ByteIndex(0));
 98           let mut index_previous = None::<usize>;
 99           terminal::indent!(writer, index_width + 3, |writer| {
100              let index = *indent_index.borrow();
101  
102              let style = if highlighted.borrow().contains(&index) {
103                 STYLE_JUMP_ADDRESS
104              } else {
105                 STYLE_GUTTER
106              };
107  
108              let index = *index;
109  
110              with(writer, style, |writer| {
111                 if index_previous == Some(index) {
112                    let dot_width = 2 + terminal::number_hex_width(index);
113                    let space_width = index_width - dot_width;
114  
115                    write!(writer, "{:>space_width$}", "")?;
116  
117                    for _ in 0..dot_width {
118                       write!(writer, "{DOT}")?;
119                    }
120                 } else {
121                    write!(writer, "{index:>#index_width$X}")?;
122                 }
123  
124                 write!(writer, " {TOP_TO_BOTTOM}")
125              })?;
126  
127              index_previous.replace(index);
128              Ok(index_width + 2)
129           });
130  
131           if **indent_index.borrow() < code.bytes.len() {
132              // DEDENT: "| "
133              terminal::dedent!(writer, 2);
134  
135              // INDENT: "┏━━━ ".
136              terminal::indent!(
137                 writer,
138                 header =
139                    const_str::concat!(RIGHT_TO_BOTTOM, LEFT_TO_RIGHT, LEFT_TO_RIGHT, LEFT_TO_RIGHT)
140                       .style(STYLE_GUTTER)
141              );
142  
143              with(writer, style::Color::Red.fg().bold(), |writer| {
144                 write!(writer, "{code_index:#X} ")
145              })?;
146  
147              match code_type {
148                 CodeType::Suspend => write(writer, &"(suspend)".cyan().bold())?,
149                 CodeType::Lambda => write(writer, &"(lambda)".magenta().bold())?,
150              }
151           }
152  
153           let mut indent: usize = 0;
154  
155           let mut items = code.iter().peekable();
156           while let Some((index, item)) = items.next() {
157              *indent_index.borrow_mut() = index;
158  
159              match item {
160                 CodeItem::Operation(operation) => {
161                    terminal::indent!(
162                       writer,
163                       (indent - usize::from(operation == Operation::ScopeEnd))
164                          * INDENT_WIDTH as usize,
165                    );
166  
167                    writeln!(writer)?;
168  
169                    if operation == Operation::ScopeEnd {
170                       indent -= 1;
171                       write(writer, &"}".style(COLORS[indent % COLORS.len()]))?;
172                       write!(writer, " ")?;
173                    }
174  
175                    with(writer, style::Color::Yellow.fg(), |writer| {
176                       write!(writer, "{operation:?}")
177                    })?;
178  
179                    if operation == Operation::ScopeStart {
180                       write!(writer, " ")?;
181                       write(writer, &"{".style(COLORS[indent % COLORS.len()]))?;
182                       indent += 1;
183                    }
184  
185                    if let Some(&(_, CodeItem::Argument(_))) = items.peek() {
186                       write(writer, &'('.bright_black().bold())?;
187                    }
188                 },
189  
190                 CodeItem::Argument(argument) => {
191                    match argument {
192                       Argument::U16(u16) => write(writer, &u16.magenta())?,
193  
194                       Argument::U64(u64) => write(writer, &u64.blue())?,
195  
196                       Argument::ValueIndex(value_index) => {
197                          let value_index_unique = code_index.add(2) * value_index.add(2);
198  
199                          with(writer, style::Color::Blue.fg().bold(), |writer| {
200                             write!(writer, "{value_index:#X} ", value_index = *value_index)
201                          })?;
202  
203                          match code[value_index] {
204                             Value::Code {
205                                is_lambda,
206                                ref code,
207                             } => {
208                                codes.push_front((
209                                   value_index_unique,
210                                   if is_lambda {
211                                      CodeType::Lambda
212                                   } else {
213                                      CodeType::Suspend
214                                   },
215                                   code,
216                                ));
217  
218                                write(writer, &"-> ".bright_black().bold())?;
219                                with(writer, style::Color::Red.fg().bold(), |writer| {
220                                   write!(writer, "{value_index_unique:#X}")
221                                })?;
222                             },
223  
224                             ref value => {
225                                // TODO: This isn't indented as our indenter:
226                                //
227                                // 1. Doesn't support dynamic indent widths. This is easily fixed by
228                                //    storing an indent (like std::iter::continuations) and measuring
229                                //    the length in IndentWrite::width.
230                                //
231                                // 2. Doesn't support specifying `Place`.
232                                //
233                                // Fix this soon, maybe --RGB.
234                                write(writer, &":: ".bright_black().bold())?;
235                                value.display_styled(writer)?;
236                             },
237                          }
238                       },
239  
240                       Argument::ByteIndex(byte_index) => {
241                          highlighted.borrow_mut().push(byte_index);
242  
243                          with(writer, STYLE_JUMP_ADDRESS, |writer| {
244                             write!(writer, "{byte_index:#X}", byte_index = *byte_index)
245                          })?;
246                       },
247                    }
248  
249                    let delimiter = match items.peek() {
250                       Some(&(_, CodeItem::Argument(_))) => ", ",
251                       _ => ")",
252                    };
253  
254                    write(writer, &delimiter.bright_black().bold())?;
255                 },
256              }
257           }
258  
259           if !codes.is_empty() {
260              writeln!(writer)?;
261              writeln!(writer)?;
262           }
263        }
264  
265        Ok(())
266     }
267  }
268  
269  #[derive(Debug, Clone, Copy, PartialEq, Eq)]
270  pub enum CodeItem {
271     Operation(Operation),
272     Argument(Argument),
273  }
274  
275  impl CodeItem {
276     pub fn as_operation(&self) -> Option<&Operation> {
277        if let &Self::Operation(ref operation) = self {
278           Some(operation)
279        } else {
280           None
281        }
282     }
283  
284     pub fn as_argument(&self) -> Option<&Argument> {
285        if let &Self::Argument(ref argument) = self {
286           Some(argument)
287        } else {
288           None
289        }
290     }
291  }
292  
293  impl Code {
294     pub fn value(&mut self, value: Value) -> ValueIndex {
295        let index = ValueIndex(self.values.len());
296        self.values.push(value);
297        index
298     }
299  
300     pub fn iter(&self) -> impl Iterator<Item = (ByteIndex, CodeItem)> {
301        gen move {
302           let mut index = ByteIndex(0);
303  
304           while *index < self.bytes.len() {
305              let (_, operation, size) = self.read_operation(index);
306  
307              yield (index, CodeItem::Operation(operation));
308              *index += size;
309  
310              match operation {
311                 Operation::Push => {
312                    let (value, size) = self.read_u64(index);
313  
314                    yield (
315                       index,
316                       CodeItem::Argument(Argument::ValueIndex(ValueIndex(
317                          usize::try_from(value).expect("value index must be valid"),
318                       ))),
319                    );
320  
321                    *index += size;
322                 },
323  
324                 Operation::Jump | Operation::JumpIf | Operation::JumpIfError => {
325                    let (value, size) = self.read_u16(index);
326  
327                    yield (
328                       index,
329                       CodeItem::Argument(Argument::ByteIndex(ByteIndex(usize::from(value)))),
330                    );
331  
332                    *index += size;
333                 },
334  
335                 Operation::Interpolate => {
336                    let (value, size) = self.read_u64(index);
337  
338                    yield (index, CodeItem::Argument(Argument::U64(value)));
339  
340                    *index += size;
341                 },
342  
343                 _ => {},
344              }
345           }
346        }
347     }
348  }
349  
350  impl ops::Index<ValueIndex> for Code {
351     type Output = Value;
352  
353     fn index(&self, index: ValueIndex) -> &Self::Output {
354        self.values.get(*index).expect("value index must be valid")
355     }
356  }
357  
358  impl Code {
359     #[must_use]
360     pub fn new(path: value::Path) -> Self {
361        Self {
362           bytes: Vec::new(),
363  
364           path,
365           spans: Vec::new(),
366  
367           values: Vec::new(),
368        }
369     }
370  
371     #[must_use]
372     pub fn path(&self) -> &value::Path {
373        &self.path
374     }
375  
376     pub fn push_u64(&mut self, data: u64) -> ByteIndex {
377        let mut encoded = [0; ENCODED_U64_LEN_MAX];
378        let len = vu128::encode_u64(&mut encoded, data);
379  
380        let index = ByteIndex(self.bytes.len());
381        self.bytes.extend_from_slice(&encoded[..len]);
382        index
383     }
384  
385     #[must_use]
386     pub fn read_u64(&self, index: ByteIndex) -> (u64, usize) {
387        let encoded = match self.bytes.get(*index..*index + ENCODED_U64_LEN_MAX) {
388           Some(slice) => {
389              <[u8; ENCODED_U64_LEN_MAX]>::try_from(slice).expect("size was statically checked")
390           },
391  
392           None => {
393              let mut buffer = [0; ENCODED_U64_LEN_MAX];
394              buffer[..self.bytes.len() - *index]
395                 .copy_from_slice(self.bytes.get(*index..).expect("byte index must be valid"));
396              buffer
397           },
398        };
399  
400        vu128::decode_u64(&encoded)
401     }
402  
403     pub fn push_u16(&mut self, data: u16) -> ByteIndex {
404        let index = ByteIndex(self.bytes.len());
405        self.bytes.extend_from_slice(&data.to_le_bytes());
406        index
407     }
408  
409     #[must_use]
410     pub fn read_u16(&self, index: ByteIndex) -> (u16, usize) {
411        let encoded = <[u8; ENCODED_U16_LEN_MAX]>::try_from(
412           self
413              .bytes
414              .get(*index..*index + ENCODED_U16_LEN_MAX)
415              .expect("byte index must be valid"),
416        )
417        .expect("size was statically checked");
418  
419        (u16::from_le_bytes(encoded), ENCODED_U16_LEN_MAX)
420     }
421  
422     pub fn push_operation(&mut self, span: Span, operation: Operation) -> ByteIndex {
423        let index = ByteIndex(self.bytes.len());
424        self.bytes.push(operation as _);
425  
426        // No need to insert the span again if this instruction was created from the
427        // last span.
428        if self
429           .spans
430           .last()
431           .is_none_or(|&(_, last_span)| last_span != span)
432        {
433           self.spans.push((index, span));
434        }
435  
436        index
437     }
438  
439     #[must_use]
440     #[track_caller]
441     pub fn read_operation(&self, index: ByteIndex) -> (value::Location, Operation, usize) {
442        let position = self.spans.partition_point(|&(index2, _)| index >= index2);
443  
444        let (_, span) = self.spans[position.saturating_sub(1)];
445  
446        (
447           value::Location::new(self.path.dupe(), span),
448           Operation::try_from(self.bytes[*index]).expect("byte index must be valid"),
449           ENCODED_OPERATION_LEN,
450        )
451     }
452  
453     pub fn point_here(&mut self, index: ByteIndex) {
454        let here = u16::try_from(self.bytes.len()).expect("bytes len must fit in u16");
455  
456        self.bytes[*index..*index + ENCODED_U16_LEN_MAX].copy_from_slice(&here.to_le_bytes());
457     }
458  }