submilli_engine/codegen/
string_pool.rs

1//! Constant-string pool — pre-codegen sweep that collects every string literal
2//! and makes each distinct value a passive data segment in the emitted module.
3
4use std::collections::BTreeMap;
5
6use crate::ExprId;
7use crate::codegen::{internal_failure, wasm_u32};
8use crate::compiler_error::CompilerFailure;
9use crate::literal_units::literal_units;
10
11#[derive(Default, Clone, Debug)]
12pub struct StringPool {
13    /// First-appearance order; index equals data-segment index.
14    pub strings: Vec<String>,
15    pub locations: BTreeMap<ExprId, usize>,
16    /// For sites without an ExprId (e.g. switch case-label strings).
17    pub text_to_idx: BTreeMap<String, usize>,
18}
19
20impl StringPool {
21    pub(crate) fn intern_text(&mut self, text: &str) -> usize {
22        if let Some(&existing) = self.text_to_idx.get(text) {
23            return existing;
24        }
25        let idx = self.strings.len();
26        self.strings.push(text.to_string());
27        self.text_to_idx.insert(text.to_string(), idx);
28        idx
29    }
30
31    pub(crate) fn record_expr(&mut self, expr_id: ExprId, text: &str) {
32        let idx = self.intern_text(text);
33        self.locations.insert(expr_id, idx);
34    }
35
36    pub fn lookup_text(&self, text: &str) -> Option<usize> {
37        self.text_to_idx.get(text).copied()
38    }
39
40    pub fn is_empty(&self) -> bool {
41        self.strings.is_empty()
42    }
43
44    /// UTF-16 code unit count for the literal at `idx` (= JS `.length`).
45    pub fn code_units(&self, idx: usize) -> Result<u32, CompilerFailure> {
46        wasm_u32(literal_units(self.string(idx)?).len())
47    }
48
49    pub fn utf16_le_bytes(&self, idx: usize) -> Result<Vec<u8>, CompilerFailure> {
50        Ok(literal_units(self.string(idx)?)
51            .into_iter()
52            .flat_map(u16::to_le_bytes)
53            .collect())
54    }
55
56    fn string(&self, idx: usize) -> Result<&str, CompilerFailure> {
57        self.strings
58            .get(idx)
59            .map(String::as_str)
60            .ok_or_else(|| internal_failure("a string literal pool index is out of range"))
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::StringPool;
67    use crate::codegen::analysis::CodegenAnalysis;
68    use crate::{Asi, Token, TokenKind, capture, check, desugar, infer, parse};
69
70    fn pool(source: &str) -> StringPool {
71        let mut asi = Asi::new(source, crate::FileId(0));
72        let mut tokens: Vec<Token> = Vec::new();
73        loop {
74            let tok = asi.next_token();
75            let is_eof = matches!(tok.kind, TokenKind::Eof);
76            tokens.push(tok);
77            if is_eof {
78                break;
79            }
80        }
81        let _ = asi.into_diagnostics();
82        let (ast, _) = parse(source, tokens, crate::FileId(0));
83        let (prelude_defs, host_defs, _) =
84            crate::runtime::prelude::cached_runtime_package_declarations();
85        let mut packages = Vec::with_capacity(prelude_defs.len() + host_defs.len());
86        packages.extend(prelude_defs.iter());
87        packages.extend(host_defs.iter());
88        let (mut ta, mut diags) = infer(source, "main", &ast, &packages);
89        diags.extend(check(&ta).unwrap());
90        ta = capture(ta).unwrap();
91        ta = desugar(ta, crate::FileId(0)).unwrap();
92        assert!(diags.is_empty(), "unexpected typecheck diags: {diags:?}");
93        CodegenAnalysis::collect(&ta, &[]).unwrap().string_pool
94    }
95
96    #[test]
97    fn empty_program_pool_is_empty() {
98        let p = pool("function main(): void { }");
99        assert!(p.is_empty());
100        assert!(p.strings.is_empty());
101        assert!(p.locations.is_empty());
102    }
103
104    #[test]
105    fn single_string_in_let_initializer() {
106        let p = pool(r#"let x: string = "hello"; function main(): void { }"#);
107        assert_eq!(p.strings, vec!["hello".to_string()]);
108        assert_eq!(p.locations.len(), 1);
109        assert_eq!(p.code_units(0).unwrap(), 5);
110        assert_eq!(p.utf16_le_bytes(0).unwrap(), b"h\0e\0l\0l\0o\0".to_vec());
111        assert!(p.code_units(1).is_err());
112        assert!(p.utf16_le_bytes(1).is_err());
113    }
114
115    #[test]
116    fn dedup_same_string_two_locations() {
117        let p =
118            pool(r#"let x: string = "hello"; let y: string = "hello"; function main(): void { }"#);
119        assert_eq!(p.strings, vec!["hello".to_string()]);
120        assert_eq!(p.locations.len(), 2);
121        for idx in p.locations.values() {
122            assert_eq!(*idx, 0);
123        }
124    }
125
126    #[test]
127    fn multiple_distinct_strings_keep_first_appearance_order() {
128        let p =
129            pool(r#"let x: string = "hello"; let y: string = "world"; function main(): void { }"#);
130        assert_eq!(p.strings, vec!["hello".to_string(), "world".to_string()]);
131        assert_eq!(p.locations.len(), 2);
132    }
133
134    #[test]
135    fn string_in_function_body_is_collected() {
136        let p = pool(r#"function main(): void { let s: string = "hi"; }"#);
137        assert_eq!(p.strings, vec!["hi".to_string()]);
138        assert_eq!(p.locations.len(), 1);
139    }
140
141    #[test]
142    fn string_in_concat_is_collected() {
143        let p = pool(r#"let z: string = "a" + "b"; function main(): void { }"#);
144        assert_eq!(p.strings, vec!["a".to_string(), "b".to_string()]);
145        assert_eq!(p.locations.len(), 2);
146    }
147
148    #[test]
149    fn string_inside_if_condition_branches_collected() {
150        let p = pool(
151            r#"
152            function main(): void {
153                const a: string = "a";
154                if (a === "b") {
155                    let x: string = "then";
156                } else {
157                    let y: string = "else";
158                }
159            }
160            "#,
161        );
162        let mut s = p.strings.clone();
163        s.sort();
164        assert_eq!(s, vec!["a", "b", "else", "then"]);
165    }
166}
submilli_engine/codegen/string_pool.rs

submilli_engine/codegen/
string_pool.rs
