git.delta.rocks / jrsonnet / refs/commits / 5ac33cfd81f9

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)]23macro_rules! const_tls {24	(const $name:ident: $t:ty = $expr:expr;) => {25		#[thread_local]26		static $name: NightlyLocalKey<$t> = NightlyLocalKey($expr);27	};28}29#[cfg(not(nightly))]30macro_rules! const_tls {31	(const $name:ident: $t:ty = $expr:expr;) => {32		thread_local! {33			static $name: $t = const { $expr };34		}35	};36}3738const_tls! {39	const STACK_LIMIT: StackLimit = StackLimit {40		max_stack_size: Cell::new(200),41		current_depth: Cell::new(0),42	};43}4445pub struct StackOverflowError;46impl From<StackOverflowError> for ErrorKind {47	fn from(_: StackOverflowError) -> Self {48		Self::StackOverflow49	}50}51impl From<StackOverflowError> for Error {52	fn from(_: StackOverflowError) -> Self {53		ErrorKind::StackOverflow.into()54	}55}5657/// Used to implement stack depth limitation58pub struct StackDepthGuard(PhantomData<()>);59impl Drop for StackDepthGuard {60	fn drop(&mut self) {61		STACK_LIMIT.with(|limit| limit.current_depth.set(limit.current_depth.get() - 1))62	}63}6465// #[cfg(feature = "nightly")]66pub fn check_depth() -> Result<StackDepthGuard, StackOverflowError> {67	STACK_LIMIT.with(|limit| {68		let current = limit.current_depth.get();69		if current < limit.max_stack_size.get() {70			limit.current_depth.set(current + 1);71			Ok(StackDepthGuard(PhantomData))72		} else {73			Err(StackOverflowError)74		}75	})76}7778pub struct StackDepthLimitOverrideGuard {79	old_limit: usize,80}81impl Drop for StackDepthLimitOverrideGuard {82	fn drop(&mut self) {83		STACK_LIMIT.with(|limit| limit.max_stack_size.set(self.old_limit));84	}85}8687pub fn limit_stack_depth(depth_limit: usize) -> StackDepthLimitOverrideGuard {88	STACK_LIMIT.with(|limit| {89		let old_limit = limit.max_stack_size.get();90		let current_depth = limit.current_depth.get();9192		limit.max_stack_size.set(current_depth + depth_limit);93		StackDepthLimitOverrideGuard { old_limit }94	})95}9697/// Like [`limit_stack_depth`], but set depth is not guarded, and will be kept98///99/// Used to implement `set_max_stack` in C api, prefer to use [`limit_stack_depth`] instead100pub fn set_stack_depth_limit(depth_limit: usize) {101	std::mem::forget(limit_stack_depth(depth_limit));102}