rsjsonnet_lang/
span.rs

1use std::num::NonZeroU64;
2
3use crate::FHashMap;
4
5#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
6pub struct SpanId(NonZeroU64);
7
8impl SpanId {
9    const OFFSET_BITS: u32 = 38;
10    const OFFSET_MASK: u64 = (1 << Self::OFFSET_BITS) - 1;
11    const LEN_MAX: u64 = (1 << (63 - Self::OFFSET_BITS)) - 1;
12
13    fn expand(self) -> ExpandedSpanId {
14        let inner = self.0.get();
15        if (inner & (1 << 63)) == 0 {
16            let offset = (inner & Self::OFFSET_MASK) - 1;
17            let len = (inner >> Self::OFFSET_BITS) as usize;
18            ExpandedSpanId::Inline(offset, len)
19        } else {
20            ExpandedSpanId::Interned((inner & !(1 << 63)) as usize)
21        }
22    }
23}
24
25impl std::fmt::Debug for SpanId {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        match self.expand() {
28            ExpandedSpanId::Inline(offset, len) => f
29                .debug_struct("Inline")
30                .field("offset", &offset)
31                .field("len", &len)
32                .finish(),
33            ExpandedSpanId::Interned(i) => f.debug_tuple("Interned").field(&i).finish(),
34        }
35    }
36}
37
38enum ExpandedSpanId {
39    Inline(u64, usize),
40    Interned(usize),
41}
42
43#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
44pub struct SpanContextId(usize);
45
46#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
47pub struct SourceId(usize);
48
49#[derive(Clone, Debug, PartialEq, Eq)]
50pub enum SpanContext {
51    Source(SourceId),
52}
53
54pub struct SpanManager {
55    contexts: Vec<(u64, SpanContext)>,
56    sources: Vec<SpanContextId>,
57    // span interner
58    span_to_idx: FHashMap<(SpanContextId, usize, usize), usize>,
59    idx_to_span: Vec<(SpanContextId, usize, usize)>,
60}
61
62impl Default for SpanManager {
63    #[inline]
64    fn default() -> Self {
65        Self::new()
66    }
67}
68
69impl SpanManager {
70    pub fn new() -> Self {
71        Self {
72            contexts: Vec::new(),
73            sources: Vec::new(),
74            span_to_idx: FHashMap::default(),
75            idx_to_span: Vec::new(),
76        }
77    }
78
79    pub fn insert_source_context(&mut self, len: usize) -> (SpanContextId, SourceId) {
80        let source_id = SourceId(self.sources.len());
81        let context_id = self.insert_context(len, SpanContext::Source(source_id));
82        self.sources.push(context_id);
83        (context_id, source_id)
84    }
85
86    fn insert_context(&mut self, len: usize, context: SpanContext) -> SpanContextId {
87        let len = u64::try_from(len).unwrap();
88        let base_offset = if let Some(&(last_end, _)) = self.contexts.last() {
89            last_end
90        } else {
91            0
92        };
93        let context_id = SpanContextId(self.contexts.len());
94        self.contexts.push((base_offset + len + 1, context));
95        context_id
96    }
97
98    #[must_use]
99    pub fn get_context(&self, context: SpanContextId) -> &SpanContext {
100        &self.contexts[context.0].1
101    }
102
103    #[must_use]
104    fn get_context_offsets(&self, context: SpanContextId) -> (u64, u64) {
105        let i = context.0;
106        if i == 0 {
107            (0, self.contexts[0].0)
108        } else {
109            (self.contexts[i - 1].0, self.contexts[i].0)
110        }
111    }
112
113    #[must_use]
114    fn get_context_from_offset(&self, offset: u64) -> SpanContextId {
115        match self.contexts.binary_search_by_key(&offset, |entry| entry.0) {
116            Ok(i) => SpanContextId(i + 1),
117            Err(i) => SpanContextId(i),
118        }
119    }
120
121    pub fn intern_span(&mut self, context: SpanContextId, start: usize, end: usize) -> SpanId {
122        let start_u64 = u64::try_from(start).unwrap();
123        let end_u64 = u64::try_from(end).unwrap();
124        let (min_offset, max_offset) = self.get_context_offsets(context);
125        assert!(start_u64 <= end_u64);
126        let start_offset = min_offset + start_u64;
127        let end_offset = min_offset + end_u64;
128        assert!(start_offset < max_offset);
129        assert!(end_offset < max_offset);
130
131        let len = end_u64 - start_u64;
132        if len > SpanId::LEN_MAX || start_offset >= SpanId::OFFSET_MASK {
133            let span = (context, start, end);
134            let i = match self.span_to_idx.entry(span) {
135                std::collections::hash_map::Entry::Occupied(entry) => *entry.get(),
136                std::collections::hash_map::Entry::Vacant(entry) => {
137                    let i = self.idx_to_span.len();
138                    self.idx_to_span.push(span);
139                    entry.insert(i);
140                    i
141                }
142            };
143            SpanId(NonZeroU64::new((i as u64) | (1 << 63)).unwrap())
144        } else {
145            SpanId(NonZeroU64::new((start_offset + 1) | (len << SpanId::OFFSET_BITS)).unwrap())
146        }
147    }
148
149    #[must_use]
150    pub fn get_span(&self, span: SpanId) -> (SpanContextId, usize, usize) {
151        match span.expand() {
152            ExpandedSpanId::Inline(start_offset, len) => {
153                let context = self.get_context_from_offset(start_offset);
154                let (min_offset, _) = self.get_context_offsets(context);
155                let start = (start_offset - min_offset) as usize;
156                (context, start, start + len)
157            }
158            ExpandedSpanId::Interned(i) => self.idx_to_span[i],
159        }
160    }
161
162    pub(crate) fn make_surrounding_span(&mut self, start_span: SpanId, end_span: SpanId) -> SpanId {
163        let (start_ctx_id, start_start_pos, _) = self.get_span(start_span);
164        let (end_ctx_id, _, end_end_pos) = self.get_span(end_span);
165        assert_eq!(start_ctx_id, end_ctx_id);
166        assert!(start_start_pos <= end_end_pos);
167        self.intern_span(start_ctx_id, start_start_pos, end_end_pos)
168    }
169}
rsjsonnet_lang/span.rs

rsjsonnet_lang/
span.rs
