Skip to content

Commit

Permalink
perf(semantic): remove bounds checks on unresolved references stack (#…
Browse files Browse the repository at this point in the history
…4390)

`SemanticBuilder` keeps a stack of hashmaps to track unresolved references. This stack only grows, never shrinks. Use that invariant to remove bounds checks.
  • Loading branch information
overlookmotel committed Jul 21, 2024
1 parent f7da22d commit da13d93
Show file tree
Hide file tree
Showing 5 changed files with 119 additions and 38 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

11 changes: 6 additions & 5 deletions crates/oxc_semantic/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,12 @@ oxc_diagnostics = { workspace = true }
oxc_index = { workspace = true }
oxc_allocator = { workspace = true }

indexmap = { workspace = true }
phf = { workspace = true, features = ["macros"] }
rustc-hash = { workspace = true }
serde = { workspace = true, features = ["derive"], optional = true }
itertools = { workspace = true }
assert-unchecked = { workspace = true }
indexmap = { workspace = true }
phf = { workspace = true, features = ["macros"] }
rustc-hash = { workspace = true }
serde = { workspace = true, features = ["derive"], optional = true }
itertools = { workspace = true }

tsify = { workspace = true, optional = true }
wasm-bindgen = { workspace = true, optional = true }
Expand Down
48 changes: 15 additions & 33 deletions crates/oxc_semantic/src/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,9 @@ use crate::{
module_record::ModuleRecordBuilder,
node::{AstNodeId, AstNodes, NodeFlags},
reference::{Reference, ReferenceFlag, ReferenceId},
scope::{ScopeFlags, ScopeId, ScopeTree, UnresolvedReferences},
scope::{ScopeFlags, ScopeId, ScopeTree},
symbol::{SymbolFlags, SymbolId, SymbolTable},
unresolved_stack::UnresolvedReferencesStack,
JSDocFinder, Semantic,
};

Expand Down Expand Up @@ -56,9 +57,6 @@ pub struct SemanticBuilder<'a> {
pub current_node_flags: NodeFlags,
pub current_symbol_flags: SymbolFlags,
pub current_scope_id: ScopeId,
/// Current scope depth.
/// 0 is global scope. 1 is `Program`. Incremented on entering a scope, and decremented on exit.
pub current_scope_depth: usize,
/// Stores current `AstKind::Function` and `AstKind::ArrowFunctionExpression` during AST visit
pub function_stack: Vec<AstNodeId>,
// To make a namespace/module value like
Expand All @@ -73,11 +71,7 @@ pub struct SemanticBuilder<'a> {
pub scope: ScopeTree,
pub symbols: SymbolTable,

// Stack used to accumulate unresolved refs while traversing scopes.
// Indexed by scope depth. We recycle `UnresolvedReferences` instances during traversal
// to reduce allocations, so the stack grows to maximum scope depth, but never shrinks.
// See: <https://github.com/oxc-project/oxc/issues/4169>
unresolved_references: Vec<UnresolvedReferences>,
unresolved_references: UnresolvedReferencesStack,

pub(crate) module_record: Arc<ModuleRecord>,

Expand Down Expand Up @@ -105,13 +99,6 @@ impl<'a> SemanticBuilder<'a> {
let scope = ScopeTree::default();
let current_scope_id = scope.root_scope_id();

// Most programs will have at least 1 place where scope depth reaches 16,
// so initialize `unresolved_references` with this length, to reduce reallocations as it grows.
// This is just an estimate of a good initial size, but certainly better than
// `Vec`'s default initial capacity of 4.
let mut unresolved_references = vec![];
unresolved_references.resize_with(16, Default::default);

let trivias = Trivias::default();
Self {
source_text,
Expand All @@ -123,13 +110,12 @@ impl<'a> SemanticBuilder<'a> {
current_symbol_flags: SymbolFlags::empty(),
current_reference_flag: ReferenceFlag::empty(),
current_scope_id,
current_scope_depth: 1, // Start with depth for `Program`
function_stack: vec![],
namespace_stack: vec![],
nodes: AstNodes::default(),
scope,
symbols: SymbolTable::default(),
unresolved_references,
unresolved_references: UnresolvedReferencesStack::new(),
module_record: Arc::new(ModuleRecord::default()),
label_builder: LabelBuilder::default(),
build_jsdoc: false,
Expand Down Expand Up @@ -200,9 +186,8 @@ impl<'a> SemanticBuilder<'a> {
}
}

debug_assert_eq!(self.current_scope_depth, 1);
self.scope.root_unresolved_references =
self.unresolved_references.into_iter().next().unwrap();
debug_assert_eq!(self.unresolved_references.scope_depth(), 1);
self.scope.root_unresolved_references = self.unresolved_references.into_root();

let jsdoc = if self.build_jsdoc { self.jsdoc.build() } else { JSDocFinder::default() };

Expand Down Expand Up @@ -356,7 +341,8 @@ impl<'a> SemanticBuilder<'a> {
let reference_flag = *reference.flag();
let reference_id = self.symbols.create_reference(reference);

self.unresolved_references[self.current_scope_depth]
self.unresolved_references
.current_mut()
.entry(reference_name)
.or_default()
.push((reference_id, reference_flag));
Expand All @@ -381,10 +367,7 @@ impl<'a> SemanticBuilder<'a> {
}

fn resolve_references_for_current_scope(&mut self) {
// `iter_mut` to get mut references to 2 entries of `unresolved_references` simultaneously
let mut iter = self.unresolved_references.iter_mut();
let parent_refs = iter.nth(self.current_scope_depth - 1).unwrap();
let current_refs = iter.next().unwrap();
let (current_refs, parent_refs) = self.unresolved_references.current_and_parent_mut();

for (name, mut references) in current_refs.drain() {
// Try to resolve a reference.
Expand Down Expand Up @@ -481,12 +464,7 @@ impl<'a> Visit<'a> for SemanticBuilder<'a> {
self.scope.get_bindings_mut(self.current_scope_id).extend(bindings);
}

// Increment scope depth, and ensure stack is large enough that
// `self.unresolved_references[self.current_scope_depth]` is initialized
self.current_scope_depth += 1;
if self.unresolved_references.len() <= self.current_scope_depth {
self.unresolved_references.push(UnresolvedReferences::default());
}
self.unresolved_references.increment_scope_depth();
}

// NB: Not called for `Program`
Expand All @@ -502,7 +480,7 @@ impl<'a> Visit<'a> for SemanticBuilder<'a> {
self.current_scope_id = parent_id;
}

self.current_scope_depth -= 1;
self.unresolved_references.decrement_scope_depth();
}

// Setup all the context for the binder.
Expand Down Expand Up @@ -554,6 +532,8 @@ impl<'a> Visit<'a> for SemanticBuilder<'a> {
}
self.current_scope_id = self.scope.add_root_scope(self.current_node_id, flags);
program.scope_id.set(Some(self.current_scope_id));
// NB: Don't call `self.unresolved_references.increment_scope_depth()`
// as scope depth is initialized as 1 already (the scope depth for `Program`).

if let Some(hashbang) = &program.hashbang {
self.visit_hashbang(hashbang);
Expand All @@ -572,6 +552,8 @@ impl<'a> Visit<'a> for SemanticBuilder<'a> {
// Don't call `leave_scope` here as `Program` is a special case - scope has no `parent_id`.
// This simplifies `leave_scope`.
self.resolve_references_for_current_scope();
// NB: Don't call `self.unresolved_references.decrement_scope_depth()`
// as scope depth must remain >= 1.

self.leave_node(kind);
}
Expand Down
1 change: 1 addition & 0 deletions crates/oxc_semantic/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ mod node;
mod reference;
mod scope;
mod symbol;
mod unresolved_stack;

pub mod dot;

Expand Down
96 changes: 96 additions & 0 deletions crates/oxc_semantic/src/unresolved_stack.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
use assert_unchecked::assert_unchecked;

use crate::scope::UnresolvedReferences;

// Stack used to accumulate unresolved refs while traversing scopes.
// Indexed by scope depth. We recycle `UnresolvedReferences` instances during traversal
// to reduce allocations, so the stack grows to maximum scope depth, but never shrinks.
// See: <https://github.com/oxc-project/oxc/issues/4169>
// This stack abstraction uses the invariant that stack only grows to avoid bounds checks.
pub(crate) struct UnresolvedReferencesStack {
stack: Vec<UnresolvedReferences>,
/// Current scope depth.
/// 0 is global scope. 1 is `Program`.
/// Incremented on entering a scope, and decremented on exit.
current_scope_depth: usize,
}

impl UnresolvedReferencesStack {
// Most programs will have at least 1 place where scope depth reaches 16,
// so initialize `stack` with this length, to reduce reallocations as it grows.
// This is just an estimate of a good initial size, but certainly better than
// `Vec`'s default initial capacity of 4.
// SAFETY: Must be >= 2 to ensure soundness of `current_and_parent_mut`.
const INITIAL_SIZE: usize = 16;
// Initial scope depth.
// Start on 1 (`Program` scope depth).
// SAFETY: Must be >= 1 to ensure soundness of `current_and_parent_mut`.
const INITIAL_DEPTH: usize = 1;
#[allow(clippy::assertions_on_constants)]
const _SIZE_CHECK: () = assert!(Self::INITIAL_SIZE > Self::INITIAL_DEPTH);

pub(crate) fn new() -> Self {
let mut stack = vec![];
stack.resize_with(Self::INITIAL_SIZE, UnresolvedReferences::default);
Self { stack, current_scope_depth: Self::INITIAL_DEPTH }
}

pub(crate) fn increment_scope_depth(&mut self) {
self.current_scope_depth += 1;

// Grow stack if required to ensure `self.stack[self.depth]` is in bounds
if self.stack.len() <= self.current_scope_depth {
self.stack.push(UnresolvedReferences::default());
}
}

pub(crate) fn decrement_scope_depth(&mut self) {
self.current_scope_depth -= 1;
// This assertion is required to ensure depth does not go below 1.
// If it did, would cause UB in `current_and_parent_mut`, which relies on
// `current_scope_depth - 1` always being a valid index into `self.stack`.
assert!(self.current_scope_depth > 0);
}

pub(crate) fn scope_depth(&self) -> usize {
self.current_scope_depth
}

/// Get unresolved references hash map for current scope
pub(crate) fn current_mut(&mut self) -> &mut UnresolvedReferences {
// SAFETY: `stack.len() > current_scope_depth` initially.
// Thereafter, `stack` never shrinks, only grows.
// `current_scope_depth` is only increased in `increment_scope_depth`,
// and it grows `stack` to ensure `stack.len()` always exceeds `current_scope_depth`.
// So this read is always guaranteed to be in bounds.
unsafe { self.stack.get_unchecked_mut(self.current_scope_depth) }
}

/// Get unresolved references hash maps for current scope, and parent scope
pub(crate) fn current_and_parent_mut(
&mut self,
) -> (&mut UnresolvedReferences, &mut UnresolvedReferences) {
// Assert invariants to remove bounds checks in code below.
// https://godbolt.org/z/vv5Wo5csv
// SAFETY: `current_scope_depth` starts at 1, and is only decremented
// in `decrement_scope_depth` which checks it doesn't go below 1.
unsafe { assert_unchecked!(self.current_scope_depth > 0) };
// SAFETY: `stack.len() > current_scope_depth` initially.
// Thereafter, `stack` never shrinks, only grows.
// `current_scope_depth` is only increased in `increment_scope_depth`,
// and it grows `stack` to ensure `stack.len()` always exceeds `current_scope_depth`.
unsafe { assert_unchecked!(self.stack.len() > self.current_scope_depth) };

let mut iter = self.stack.iter_mut();
let parent = iter.nth(self.current_scope_depth - 1).unwrap();
let current = iter.next().unwrap();
(current, parent)
}

pub(crate) fn into_root(self) -> UnresolvedReferences {
// SAFETY: Stack starts with a non-zero size and never shrinks.
// This assertion removes bounds check in `.next()`.
unsafe { assert_unchecked!(!self.stack.is_empty()) };
self.stack.into_iter().next().unwrap()
}
}

0 comments on commit da13d93

Please sign in to comment.