Skip to main content

aranya_policy_module/
codemap.rs

1extern crate alloc;
2
3use alloc::{string::String, vec::Vec};
4use core::cmp::Ordering;
5
6use aranya_policy_ast::Span;
7use serde::{Deserialize, Serialize};
8
9/// An error for a range that doesn't exist. Used in [CodeMap].
10#[derive(Debug, thiserror::Error)]
11#[error("range error")]
12pub struct RangeError;
13
14/// This is a simplified version of Pest's `Span`. We can't use Pest's version because we
15/// need to work in `no_std` environments.
16pub struct SpannedText<'a> {
17    text: &'a str,
18    start: usize,
19    end: usize,
20}
21
22impl<'a> SpannedText<'a> {
23    /// Create a span inside a text reference. `start` and `end` are expressed in bytes. If
24    /// the `start` or `end` do not occur on a UTF-8 character boundary, this will return
25    /// `None`.
26    pub fn new(text: &'a str, start: usize, end: usize) -> Option<Self> {
27        if text.get(start..end).is_some() {
28            Some(SpannedText { text, start, end })
29        } else {
30            None
31        }
32    }
33
34    /// The start of the span, in bytes. This is guaranteed to be on a UTF-8 boundary.
35    pub fn start(&self) -> usize {
36        self.start
37    }
38
39    /// The end of the span, in bytes. This is guaranteed to be on a UTF-8 boundary.
40    pub fn end(&self) -> usize {
41        self.end
42    }
43
44    /// Return the span as a &str. Assumes the start and end positions
45    /// are character-aligned.
46    pub fn as_str(&self) -> &str {
47        &self.text[self.start..self.end]
48    }
49
50    /// Calculate the line and column position, in characters.
51    fn linecol(&self, pos: usize) -> (usize, usize) {
52        assert!(pos < self.text.len());
53        let mut line: usize = 1;
54        let mut col: usize = 1;
55        for c in self.text[0..pos].chars() {
56            if c == '\n' {
57                line = line.checked_add(1).expect("line + 1 must not wrap");
58                col = 1;
59            } else {
60                col = col.checked_add(1).expect("col + 1 must not wrap");
61            }
62        }
63
64        (line, col)
65    }
66
67    /// Returns the line and column of the start position.
68    pub fn start_linecol(&self) -> (usize, usize) {
69        self.linecol(self.start)
70    }
71
72    /// Returns the line and column of the end position.
73    pub fn end_linecol(&self) -> (usize, usize) {
74        self.linecol(self.end)
75    }
76}
77
78/// The code map contains the original source and can map VM instructions to text ranges
79/// inside that source.
80#[derive(
81    Debug,
82    Clone,
83    Eq,
84    PartialEq,
85    Serialize,
86    Deserialize,
87    rkyv::Archive,
88    rkyv::Deserialize,
89    rkyv::Serialize,
90)]
91pub struct CodeMap {
92    /// The original policy source code
93    text: String,
94    /// A mapping between ranges of instructions and source spans.
95    mapping: Vec<(usize, Span)>,
96}
97
98impl CodeMap {
99    /// Create a new, empty CodeMap from a text and a set of ranges.
100    pub fn new(text: impl Into<String>) -> Self {
101        Self {
102            text: text.into(),
103            mapping: Vec::new(),
104        }
105    }
106
107    /// Get the original source code.
108    pub fn text(&self) -> &str {
109        &self.text
110    }
111
112    /// Add new mapping starting at `instruction` which maps to `span`.
113    pub fn map_instruction(&mut self, instruction: usize, span: Span) -> Result<(), RangeError> {
114        match self
115            .mapping
116            .last_mut()
117            .map(|last| (instruction.cmp(&last.0), last))
118        {
119            // Don't break sorting.
120            Some((Ordering::Less, _)) => return Err(RangeError),
121            // Update existing span to be more specific.
122            Some((Ordering::Equal, last)) => last.1 = span,
123            // Add new span to end.
124            Some((Ordering::Greater, _)) | None => self.mapping.push((instruction, span)),
125        }
126        Ok(())
127    }
128
129    /// Retrieve the [`Span`] containing the given instruction pointer.
130    pub fn span_from_instruction(&self, ip: usize) -> Result<SpannedText<'_>, RangeError> {
131        let idx = match self.mapping.binary_search_by(|(i, _)| i.cmp(&ip)) {
132            Ok(idx) => idx,
133            Err(idx) => idx.checked_sub(1).ok_or(RangeError)?,
134        };
135        let span = self.mapping[idx].1;
136        SpannedText::new(&self.text, span.start(), span.end()).ok_or(RangeError)
137    }
138}