Allocate proc bodies in a separate mimalloc heap

crates/dm-langserver/src/jrpc_io.rs

@@ -3,6 +3,7 @@
 //! JSON-RPC over stdin/stdout with Content-Length headers.
 
 use std::io::{self, BufRead, Write};
+use std::sync::mpsc;
 
 pub fn run_until_stdin_eof<F: FnMut(&str)>(mut f: F) {
     let stdin = io::stdin();
@@ -12,6 +13,31 @@ pub fn run_until_stdin_eof<F: FnMut(&str)>(mut f: F) {
     }
 }
 
+pub enum Event {
+    /// A message from stdin, or `None` at EOF.
+    Read(Result<Option<String>, String>),
+    /// A background job finished.
+    JobFinished,
+}
+
+/// Read stdin on its own thread, so background threads can also send events.
+pub fn stdin_events() -> (mpsc::Sender<Event>, mpsc::Receiver<Event>) {
+    let (tx, rx) = mpsc::channel();
+    let reader_tx = tx.clone();
+    std::thread::spawn(move || {
+        let mut stdin = io::stdin().lock();
+        loop {
+            let read = read(&mut stdin).map_err(|error| error.to_string());
+            let done = !matches!(read, Ok(Some(_)));
+            let _ = reader_tx.send(Event::Read(read));
+            if done {
+                break;
+            }
+        }
+    });
+    (tx, rx)
+}
+
 pub fn run_with_read<R: BufRead, F: FnMut(&str)>(input: &mut R, mut f: F) {
     while let Some(message) = read(input).expect("JSON-RPC read error") {
         f(&message);

crates/dm-langserver/src/main.rs

@@ -26,7 +26,7 @@ use std::collections::VecDeque;
 use std::collections::hash_map::Entry;
 use std::path::PathBuf;
 use std::rc::Rc;
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, Mutex, mpsc};
 
 use clap::Parser;
 use dm::annotation::{Annotation, AnnotationTree};
@@ -48,6 +48,7 @@ mod document;
 mod extras;
 mod find_references;
 mod jrpc_io;
+mod proc_heap;
 mod symbol_search;
 
 use extras::{QueryObjectTree, Reparse, SetTraceVsc, StartDebugger};
@@ -131,8 +132,17 @@ impl DmLangserverCli {
         match &self.command {
             Some(Command::Debugger(args)) => debugger::debugger_main(args),
             None => {
-                let mut engine = Engine::new();
-                jrpc_io::run_until_stdin_eof(|message| engine.handle_input(message));
+                let (events_tx, events) = jrpc_io::stdin_events();
+                let mut engine = Engine::new(events_tx);
+                for event in events {
+                    match event {
+                        jrpc_io::Event::Read(read) => match read.expect("JSON-RPC read error") {
+                            Some(message) => engine.handle_input(&message),
+                            None => break,
+                        },
+                        jrpc_io::Event::JobFinished => engine.release_proc_bodies(),
+                    }
+                }
                 engine.exit(0);
             },
         }
@@ -316,6 +326,8 @@ struct Engine {
     context: dm::Context,
     defines: Option<dm::preprocessor::DefineHistory>,
     objtree: Arc<dm::objtree::ObjectTree>,
+    /// Lets background jobs tell the main loop when they finish.
+    events_tx: mpsc::Sender<jrpc_io::Event>,
     references_table: background::Background<find_references::ReferencesTable>,
 
     annotations: HashMap<Url, (FileId, FileId, Rc<AnnotationTree>)>,
@@ -327,7 +339,7 @@ struct Engine {
 }
 
 impl Engine {
-    fn new() -> Self {
+    fn new(events_tx: mpsc::Sender<jrpc_io::Event>) -> Self {
         Engine {
             docs: Default::default(),
 
@@ -340,6 +352,7 @@ impl Engine {
             context: dm::Context::default(),
             defines: None,
             objtree: Default::default(),
+            events_tx,
             references_table: Default::default(),
 
             annotations: Default::default(),
@@ -608,6 +621,7 @@ impl Engine {
         {
             let mut parser = dm::Parser::new(&self.context, &mut pp);
             parser.enable_procs();
+            parser.set_proc_body_scope(proc_heap::enter, proc_heap::exit);
             let (fatal_errored_2, objtree) = parser.parse_object_tree_2();
             fatal_errored = fatal_errored_2;
             self.objtree = Arc::new(objtree);
@@ -628,6 +642,7 @@ impl Engine {
 
         // Background thread: prepare the Find All References database.
         let references_objtree = self.objtree.clone();
+        let events_tx = self.events_tx.clone();
         self.references_table.spawn(move || {
             let table = find_references::ReferencesTable::new(&references_objtree);
             let elapsed = start.elapsed();
@@ -637,7 +652,9 @@ impl Engine {
                 elapsed.subsec_millis()
             );
             print_thread_total();
+            drop(references_objtree);
             collect_freed_memory();
+            let _ = events_tx.send(jrpc_io::Event::JobFinished);
             table
         });
 
@@ -652,6 +669,7 @@ impl Engine {
             let root = self.root.clone();
             let related_info = self.client_caps.related_info;
             let diagnostics_tracker = self.diagnostics_tracker.clone();
+            let events_tx = self.events_tx.clone();
             std::thread::spawn(move || {
                 dreamchecker::run(&context, &objtree);
                 let elapsed = start.elapsed();
@@ -672,8 +690,9 @@ impl Engine {
                 diagnostics_tracker.lock().unwrap().send(map);
 
                 issue_notification::<extras::WindowStatus>(Default::default());
-                drop(context);
+                drop((context, objtree));
                 collect_freed_memory();
+                let _ = events_tx.send(jrpc_io::Event::JobFinished);
             });
         } else {
             self.issue_notification::<extras::WindowStatus>(Default::default());
@@ -713,10 +732,6 @@ impl Engine {
             );
         }
 
-        /*if let Some(objtree) = Arc::get_mut(&mut self.objtree) {
-            objtree.drop_code();
-        }*/
-
         // Print the total time.
         print_thread_total();
 
@@ -1282,6 +1297,16 @@ impl Engine {
         jrpc_io::write(&serde_json::to_string(&response).expect("response bad to_string"));
     }
 
+    /// Drop proc bodies once the background jobs are done with them.
+    fn release_proc_bodies(&mut self) {
+        if let Some(objtree) = Arc::get_mut(&mut self.objtree) {
+            objtree.drop_code();
+            proc_heap::collect();
+            // the heaps share segments, so the main heap needs collecting too
+            collect_freed_memory();
+        }
+    }
+
     fn handle_call(&mut self, call: Call) -> Option<Output> {
         match call {
             Call::Invalid { id } => Some(Output::invalid_request(id, JSONRPC_VERSION)),

crates/dm-langserver/src/proc_heap.rs

@@ -0,0 +1,29 @@
+//! Proc bodies are only needed until dreamchecker and Find All References are done with them. Keeping them in their own heap means their memory can actually be given back once they're dropped.
+#![allow(unsafe_code)]
+
+use libmimalloc_sys::{
+    mi_heap_collect, mi_heap_get_backing, mi_heap_new, mi_heap_set_default, mi_heap_t,
+};
+
+thread_local! {
+    // SAFETY: mi_heap_new is a ffi call, no safety contract
+    static PROC_HEAP: *mut mi_heap_t = unsafe { mi_heap_new() };
+}
+
+/// Start allocating in the proc heap.
+pub fn enter() {
+    // SAFETY: the heap belongs to this thread
+    PROC_HEAP.with(|&heap| unsafe { mi_heap_set_default(heap) });
+}
+
+/// Go back to allocating in this thread's normal heap.
+pub fn exit() {
+    // SAFETY: the backing heap belongs to this thread
+    unsafe { mi_heap_set_default(mi_heap_get_backing()) };
+}
+
+/// Give the memory of dropped proc bodies back to the OS.
+pub fn collect() {
+    // SAFETY: the heap belongs to this thread
+    PROC_HEAP.with(|&heap| unsafe { mi_heap_collect(heap, true) });
+}

crates/dreammaker/src/parser.rs

@@ -336,6 +336,8 @@ pub struct Parser<'ctx, 'an, 'inp> {
     procs: bool,
     procs_bad: u64,
     procs_good: u64,
+    proc_body_enter: fn(),
+    proc_body_exit: fn(),
 }
 
 impl<'ctx, 'an, 'inp> HasLocation for Parser<'ctx, 'an, 'inp> {
@@ -385,13 +387,21 @@ impl<'ctx, 'an, 'inp> Parser<'ctx, 'an, 'inp> {
             procs: false,
             procs_bad: 0,
             procs_good: 0,
+            proc_body_enter: || {},
+            proc_body_exit: || {},
         }
     }
 
     pub fn enable_procs(&mut self) {
         self.procs = true;
     }
 
+    /// Run `enter` before and `exit` after parsing each proc body so we can keep proc bodies in their own heap.
+    pub fn set_proc_body_scope(&mut self, enter: fn(), exit: fn()) {
+        self.proc_body_enter = enter;
+        self.proc_body_exit = exit;
+    }
+
     pub fn annotate_to(&mut self, annotations: &'an mut AnnotationTree) {
         self.annotations = Some(annotations);
         self.procs = true;
@@ -1270,6 +1280,8 @@ impl<'ctx, 'an, 'inp> Parser<'ctx, 'an, 'inp> {
         }
 
         let code = if self.procs {
+            (self.proc_body_enter)();
+            let outer_errors = std::mem::take(&mut *self.context.errors_mut());
             let result = {
                 let mut subparser: Parser<'ctx, '_, '_> = Parser::new(self.context, body_tt);
                 if let Some(a) = self.annotations.as_mut() {
@@ -1283,13 +1295,20 @@ impl<'ctx, 'an, 'inp> Parser<'ctx, 'an, 'inp> {
             } else {
                 self.procs_bad += 1;
             }
-            match result {
+            let code = match result {
                 Err(err) => {
                     self.context.register_error(err);
                     None
                 },
                 Ok(code) => Some(code),
-            }
+            };
+            (self.proc_body_exit)();
+            // errors outlive the proc body, so copy them out of its scope
+            let body_errors = std::mem::replace(&mut *self.context.errors_mut(), outer_errors);
+            self.context
+                .errors_mut()
+                .extend(body_errors.iter().cloned());
+            code
         } else {
             None
         };

crates/dreammaker/tests/ast_tests.rs

@@ -374,3 +374,23 @@ fn for_key_value_with_as() {
         },
     );
 }
+
+#[test]
+fn proc_body_errors_keep_their_order() {
+    // one diagnostic per line, from both proc parameters and proc bodies, including a body that fails to parse
+    with_code(
+        "
+/proc/one(var/a)
+    var/tmp/x = 1
+    src.
+/proc/two(var/b)
+    var/z = 1 +
+/proc/three(var/c)
+    var/tmp/y = 1
+    ",
+        |context, _| {
+            let lines: Vec<u32> = context.errors().iter().map(|e| e.location().line).collect();
+            assert_eq!(lines, [1, 2, 3, 4, 5, 6, 7]);
+        },
+    );
+}