cu29_intern_strs/
lib.rs

1use anyhow::{Context, Result};
2use byteorder::{ByteOrder, LittleEndian};
3use rkv::backend::{Lmdb, LmdbDatabase, LmdbEnvironment, LmdbRwTransaction};
4use rkv::{MultiStore, Rkv, SingleStore, StoreOptions, Value, Writer};
5use std::fs;
6use std::path::{Path, PathBuf};
7use std::sync::Mutex;
8use std::sync::OnceLock;
9
10type SStore = SingleStore<LmdbDatabase>;
11type MStore = MultiStore<LmdbDatabase>;
12type IndexType = u32;
13
14/// The name of the directory where the log index is stored.
15const INDEX_DIR_NAME: &str = "cu29_log_index";
16
17fn parent_n_times(path: &Path, n: usize) -> Option<PathBuf> {
18    let mut result = Some(path.to_path_buf());
19    for _ in 0..n {
20        result = result?.parent().map(PathBuf::from);
21    }
22    result
23}
24
25/// Convenience function to returns the default path for the log index directory.
26pub fn default_log_index_dir() -> PathBuf {
27    let outdir = std::env::var("LOG_INDEX_DIR").expect("no LOG_INDEX_DIR system variable set, be sure build.rs sets it, see cu29_log/build.rs for example.");
28    let outdir_path = Path::new(&outdir);
29    parent_n_times(outdir_path, 3).unwrap().join(INDEX_DIR_NAME)
30}
31
32/// Reads all interned strings from the index at the specified path.
33/// The index is created at compile time within your project output directory.
34pub fn read_interned_strings(index: &Path) -> Result<Vec<String>> {
35    let mut all_strings = Vec::<String>::new();
36    let env =
37        Rkv::new::<Lmdb>(index).context("Could not open the string index. Check the path.")?;
38
39    let index_to_string = env
40        .open_single("index_to_string", StoreOptions::default())
41        .context("Could not open the index_to_string store")?;
42    let db_reader = env.read().unwrap();
43    let ri = index_to_string.iter_start(&db_reader);
44    let mut i = ri.expect("Failed to start iterator");
45    while let Some(Ok(v)) = i.next() {
46        let (k, v) = v;
47        let index = LittleEndian::read_u32(k) as usize;
48
49        if let rkv::Value::Str(s) = v {
50            if all_strings.len() <= index {
51                all_strings.resize(index + 1, String::new());
52            }
53
54            all_strings[index] = s.to_string();
55        }
56    }
57    Ok(all_strings)
58}
59
60#[cfg(feature = "macro_debug")]
61const COLORED_PREFIX_BUILD_LOG: &str = "\x1b[32mCLog:\x1b[0m";
62
63#[cfg(feature = "macro_debug")]
64macro_rules! build_log {
65    ($($arg:tt)*) => {
66        eprintln!("{} {}", COLORED_PREFIX_BUILD_LOG, format!($($arg)*));
67    };
68}
69
70static RKV: OnceLock<Mutex<Rkv<LmdbEnvironment>>> = OnceLock::new();
71static DBS: OnceLock<Mutex<(SStore, SStore, SStore, MStore)>> = OnceLock::new();
72
73fn rkv() -> &'static Mutex<Rkv<LmdbEnvironment>> {
74    RKV.get_or_init(|| {
75        let target_dir = default_log_index_dir();
76
77        // Should never happen I believe.
78        if !target_dir.exists() {
79            fs::create_dir_all(&target_dir).unwrap();
80        }
81
82        #[cfg(feature = "macro_debug")]
83        {
84            build_log!(
85                "=================================================================================="
86            );
87            build_log!("Interned strings are stored in: {:?}", target_dir);
88            build_log!("   [r] is reused index and [n] is new index.");
89            build_log!(
90                "=================================================================================="
91            );
92        }
93
94        let env = Rkv::new::<Lmdb>(&target_dir).unwrap();
95        Mutex::new(env)
96    })
97}
98
99fn dbs() -> &'static Mutex<(SStore, SStore, SStore, MStore)> {
100    DBS.get_or_init(|| {
101        let env = rkv().lock().unwrap();
102
103        let counter = env.open_single("counter", StoreOptions::create()).unwrap();
104        let index_to_string = env
105            .open_single("index_to_string", StoreOptions::create())
106            .unwrap();
107        let string_to_index = env
108            .open_single("string_to_index", StoreOptions::create())
109            .unwrap();
110        let index_to_callsites = env
111            .open_multi("index_to_callsites", StoreOptions::create())
112            .unwrap();
113
114        Mutex::new((
115            counter,
116            index_to_string,
117            string_to_index,
118            index_to_callsites,
119        ))
120    })
121}
122
123pub fn intern_string(s: &str) -> Option<IndexType> {
124    let (counter_store, index_to_string, string_to_index, _) = &mut *dbs().lock().unwrap();
125    let index = {
126        let env = rkv().lock().unwrap();
127        // If this string already exists in the store, return the index
128        {
129            let reader = env.read().unwrap();
130            // check if log_string is already in the string_to_index store
131            if let Ok(Some(Value::U64(index))) = string_to_index.get(&reader, s) {
132                #[cfg(feature = "macro_debug")]
133                {
134                    build_log!("#{:0>3} [r] -> {}.", index, s);
135                }
136                return Some(index as IndexType);
137            };
138        }
139        let mut writer = env.write().unwrap();
140        let next_index = get_next_index(&mut writer, counter_store).unwrap();
141        // Insert the new string into the store
142        index_to_string
143            .put(&mut writer, next_index.to_le_bytes(), &Value::Str(s))
144            .unwrap();
145        string_to_index
146            .put(&mut writer, s, &Value::U64(next_index as u64))
147            .unwrap();
148        writer.commit().unwrap();
149        Some(next_index)
150    };
151    #[cfg(feature = "macro_debug")]
152    {
153        build_log!("#{:0>3} [n] -> {}.", index.unwrap(), s);
154    }
155    index
156}
157
158#[allow(dead_code)]
159pub fn record_callsite(filename: &str, line_number: u32) -> Option<IndexType> {
160    intern_string(format!("{filename}:{line_number}").as_str())
161}
162
163const COUNTER_KEY: &str = "__counter__";
164fn get_next_index(
165    writer: &mut Writer<LmdbRwTransaction>,
166    counter_store: &SStore,
167) -> Result<IndexType, Box<dyn std::error::Error>> {
168    let current_counter = match counter_store.get(writer, COUNTER_KEY)? {
169        Some(Value::U64(value)) => value as IndexType,
170        _ => 0,
171    };
172
173    let next_counter = current_counter + 1;
174    counter_store.put(writer, COUNTER_KEY, &Value::U64(next_counter as u64))?;
175    Ok(next_counter)
176}
cu29_intern_strs/lib.rs

cu29_intern_strs/
lib.rs
