ast: Use RwLock and once_cell:sync::Lazy for SpanMap

compiler/hash-ast/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 [dependencies]
 num-bigint = "0.4"
 replace_with = "0.1.7"
+once_cell = "1.17"
 
 hash-source = { path = "../hash-source" }
 hash-tree-def = { path = "../hash-tree-def" }

compiler/hash-ast/src/ast.rs

@@ -5,7 +5,6 @@ use std::{
     hash::Hash,
     ops::{Deref, DerefMut},
     path::PathBuf,
-    sync::OnceLock,
 };
 
 use hash_source::{
@@ -17,8 +16,10 @@ use hash_tree_def::define_tree;
 use hash_utils::{
     counter,
     index_vec::{Idx, IndexVec},
+    parking_lot::RwLock,
     smallvec::SmallVec,
 };
+use once_cell::sync::Lazy;
 use replace_with::replace_with_or_abort;
 
 counter! {
@@ -43,27 +44,28 @@ impl Idx for AstNodeId {
 /// [`AstNode<T>`] itself in order for other data structures to be able
 /// to query the [Span] of a node simply by using the [AstNodeId] of the
 /// node.
-static mut SPAN_MAP: OnceLock<IndexVec<AstNodeId, Span>> = OnceLock::new();
+static mut SPAN_MAP: Lazy<RwLock<IndexVec<AstNodeId, Span>>> =
+    Lazy::new(|| RwLock::new(IndexVec::new()));
 
 /// Utilities for working with the [`SPAN_MAP`].
-pub struct AstUtils;
+pub struct SpanMap;
 
-impl AstUtils {
-    /// Get the span of a node by [AstNodId].
+impl SpanMap {
+    /// Get the span of a node by [AstNodeId].
     pub fn span_of(id: AstNodeId) -> Span {
-        Self::span_map()[id]
+        unsafe { SPAN_MAP.read()[id] }
     }
 
-    /// Get a reference to the [`SPAN_MAP`].
-    pub fn span_map() -> &'static IndexVec<AstNodeId, Span> {
-        unsafe { SPAN_MAP.get_or_init(IndexVec::new) }
-    }
-
-    /// Get a mutable reference to the [`SPAN_MAP`]. This is only 
+    /// Get a mutable reference to the [`SPAN_MAP`]. This is only
     /// internal to the `hash-ast` crate since it creates entries
     /// in the span map when creating new AST nodes.
-    fn span_map_mut() -> &'static mut IndexVec<AstNodeId, Span> {
-        unsafe { SPAN_MAP.get_mut().unwrap() }
+    fn add_span(span: Span) -> AstNodeId {
+        unsafe { SPAN_MAP.get_mut().push(span) }
+    }
+
+    /// Update the span of a node by [AstNodeId].
+    fn update_span(id: AstNodeId, span: Span) {
+        unsafe { SPAN_MAP.get_mut()[id] = span }
     }
 }
 
@@ -87,7 +89,7 @@ impl<T> PartialEq for AstNode<T> {
 impl<T> AstNode<T> {
     /// Create a new node with a given body and location.
     pub fn new(body: T, span: Span) -> Self {
-        let id = AstUtils::span_map_mut().push(span);
+        let id = SpanMap::add_span(span);
         Self { body: Box::new(body), id }
     }
 
@@ -108,12 +110,12 @@ impl<T> AstNode<T> {
 
     /// Get the [Span] of this [AstNode].
     pub fn span(&self) -> Span {
-        AstUtils::span_map()[self.id]
+        SpanMap::span_of(self.id)
     }
 
     /// Set the [Span] of this [AstNode].
     pub fn set_span(&mut self, span: Span) {
-        AstUtils::span_map_mut()[self.id] = span;
+        SpanMap::update_span(self.id, span)
     }
 
     /// Get the [AstNodeId] of this node.
@@ -175,7 +177,7 @@ impl<'t, T> AstNodeRef<'t, T> {
 
     /// Get the [Span] of this [AstNodeRef].
     pub fn span(&self) -> Span {
-        AstUtils::span_map()[self.id]
+        SpanMap::span_of(self.id)
     }
 
     /// Get the [AstNodeId] of this [AstNodeRef].
@@ -225,7 +227,7 @@ impl<'t, T> AstNodeRefMut<'t, T> {
 
     /// Get the [Span] of this [AstNodeRefMut].
     pub fn span(&self) -> Span {
-        AstUtils::span_map()[self.id]
+        SpanMap::span_of(self.id)
     }
 
     /// Get the [AstNodeId] of this [AstNodeRefMut].

compiler/hash-ast/src/lib.rs

@@ -1,6 +1,6 @@
 //! Hash Compiler AST library file
 
-#![feature(box_into_inner, iter_intersperse, let_chains)]
+#![feature(box_into_inner, iter_intersperse, let_chains, lazy_cell)]
 
 pub mod ast;
 pub mod node_map;

compiler/hash-parser/src/lib.rs

@@ -10,10 +10,7 @@ mod source;
 use std::{env, path::PathBuf};
 
 use crossbeam_channel::{unbounded, Sender};
-use hash_ast::{
-    ast::{self, AstUtils},
-    node_map::ModuleEntry,
-};
+use hash_ast::{ast, node_map::ModuleEntry};
 use hash_lexer::Lexer;
 use hash_pipeline::{
     interface::{CompilerInterface, CompilerStage},
@@ -64,9 +61,6 @@ impl<Ctx: ParserCtxQuery> CompilerStage<Ctx> for Parser {
         let ParserCtx { workspace, pool } = &mut ctx.data();
         let current_dir = env::current_dir().map_err(|err| vec![err.into()])?;
 
-        // Initialise the map if it hasn't been initialised yet.
-        AstUtils::span_map();
-
         let mut collected_diagnostics = Vec::new();
         let (sender, receiver) = unbounded::<ParserAction>();
 

compiler/hash-tir/Cargo.toml

@@ -13,7 +13,6 @@ textwrap = "0.16"
 utility-types = "0.0.2"
 typed-builder = "0.11"
 lazy_static = "1.4"
-parking_lot = "0.12"
 
 hash-ast = { path = "../hash-ast" }
 hash-source = { path = "../hash-source" }

compiler/hash-tir/src/ast_info.rs

@@ -2,7 +2,7 @@ use std::hash::Hash;
 
 use bimap::BiMap;
 use hash_ast::ast::AstNodeId;
-use parking_lot::RwLock;
+use hash_utils::parking_lot::RwLock;
 
 use crate::{
     args::ArgId,

compiler/hash-tir/src/locations.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 
 use hash_source::location::{SourceLocation, Span};
 use hash_storage::store::SequenceStoreKey;
-use parking_lot::RwLock;
+use hash_utils::parking_lot::RwLock;
 
 use super::{
     args::{ArgId, ArgsId, PatArgId, PatArgsId},

compiler/hash-tir/src/problems.rs

@@ -1,4 +1,4 @@
-use parking_lot::RwLock;
+use hash_utils::parking_lot::RwLock;
 
 use crate::context::Context;
 

compiler/hash-tir/src/scopes.rs

@@ -9,7 +9,7 @@ use hash_storage::{
     static_single_store,
     store::{statics::StoreId, Store, TrivialSequenceStoreKey},
 };
-use parking_lot::{MappedRwLockReadGuard, MappedRwLockWriteGuard};
+use hash_utils::parking_lot::{MappedRwLockReadGuard, MappedRwLockWriteGuard};
 use textwrap::indent;
 use utility_types::omit;
 

compiler/hash-utils/src/lib.rs

@@ -27,5 +27,6 @@ pub use fxhash;
 pub use index_vec;
 pub use itertools;
 pub use log;
+pub use parking_lot;
 pub use smallvec;
 pub use thin_vec;