git.delta.rocks / jrsonnet / refs/commits / b3f009bb7b2e

1use std::{cell::Cell, marker::PhantomData};23use crate::error::{Error, ErrorKind};45struct StackLimit {6	max_stack_size: Cell<usize>,7	current_depth: Cell<usize>,8}910#[cfg(nightly)]11struct NightlyLocalKey<T>(pub T);12#[cfg(nightly)]13impl<T> NightlyLocalKey<T> {14	#[inline(always)]15	fn with<U>(&self, v: impl FnOnce(&T) -> U) -> U {16		v(&self.0)17	}18}19#[cfg(not(nightly))]20type NightlyLocalKey<T> = std::thread::LocalKey<T>;2122#[cfg(nightly)]23#[macro_export]24macro_rules! const_tls {25	(const $name:ident: $t:ty = $expr:expr;) => {26		#[thread_local]27		static $name: NightlyLocalKey<$t> = NightlyLocalKey($expr);28	};29}30#[cfg(not(nightly))]31#[macro_export]32macro_rules! const_tls {33	(const $name:ident: $t:ty = $expr:expr;) => {34		thread_local! {35			static $name: $t = const { $expr };36		}37	};38}3940const_tls! {41	const STACK_LIMIT: StackLimit = StackLimit {42		max_stack_size: Cell::new(200),43		current_depth: Cell::new(0),44	};45}4647pub struct StackOverflowError;48impl From<StackOverflowError> for ErrorKind {49	fn from(_: StackOverflowError) -> Self {50		Self::StackOverflow51	}52}53impl From<StackOverflowError> for Error {54	fn from(_: StackOverflowError) -> Self {55		ErrorKind::StackOverflow.into()56	}57}5859/// Used to implement stack depth limitation60pub struct StackDepthGuard(PhantomData<()>);61impl Drop for StackDepthGuard {62	fn drop(&mut self) {63		STACK_LIMIT.with(|limit| limit.current_depth.set(limit.current_depth.get() - 1))64	}65}6667// #[cfg(feature = "nightly")]68pub fn check_depth() -> Result<StackDepthGuard, StackOverflowError> {69	STACK_LIMIT.with(|limit| {70		let current = limit.current_depth.get();71		if current < limit.max_stack_size.get() {72			limit.current_depth.set(current + 1);73			Ok(StackDepthGuard(PhantomData))74		} else {75			Err(StackOverflowError)76		}77	})78}7980pub struct StackDepthLimitOverrideGuard {81	old_limit: usize,82}83impl Drop for StackDepthLimitOverrideGuard {84	fn drop(&mut self) {85		STACK_LIMIT.with(|limit| limit.max_stack_size.set(self.old_limit));86	}87}8889pub fn limit_stack_depth(depth_limit: usize) -> StackDepthLimitOverrideGuard {90	STACK_LIMIT.with(|limit| {91		let old_limit = limit.max_stack_size.get();92		let current_depth = limit.current_depth.get();9394		limit.max_stack_size.set(current_depth + depth_limit);95		StackDepthLimitOverrideGuard { old_limit }96	})97}9899/// Like [`limit_stack_depth`], but set depth is not guarded, and will be kept100///101/// Used to implement `set_max_stack` in C api, prefer to use [`limit_stack_depth`] instead102pub fn set_stack_depth_limit(depth_limit: usize) {103	std::mem::forget(limit_stack_depth(depth_limit));104}