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

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]15	fn with<U>(&self, v: impl FnOnce(&T) -> U) -> U {16		v(&self.0)17	}18}19#[cfg(not(nightly))]20#[allow(dead_code)]21type NightlyLocalKey<T> = std::thread::LocalKey<T>;2223#[cfg(nightly)]24#[macro_export]25macro_rules! const_tls {26	(const $name:ident: $t:ty = $expr:expr;) => {27		#[thread_local]28		static $name: NightlyLocalKey<$t> = NightlyLocalKey($expr);29	};30}31#[cfg(not(nightly))]32#[macro_export]33macro_rules! const_tls {34	(const $name:ident: $t:ty = $expr:expr;) => {35		thread_local! {36			static $name: $t = const { $expr };37		}38	};39}4041const_tls! {42	const STACK_LIMIT: StackLimit = StackLimit {43		max_stack_size: Cell::new(200),44		current_depth: Cell::new(0),45	};46}4748pub struct StackOverflowError;49impl From<StackOverflowError> for ErrorKind {50	fn from(_: StackOverflowError) -> Self {51		Self::StackOverflow52	}53}54impl From<StackOverflowError> for Error {55	fn from(_: StackOverflowError) -> Self {56		ErrorKind::StackOverflow.into()57	}58}5960/// Used to implement stack depth limitation61pub struct StackDepthGuard(PhantomData<()>);62impl Drop for StackDepthGuard {63	fn drop(&mut self) {64		STACK_LIMIT.with(|limit| limit.current_depth.set(limit.current_depth.get() - 1));65	}66}6768// #[cfg(feature = "nightly")]69pub fn check_depth() -> Result<StackDepthGuard, StackOverflowError> {70	STACK_LIMIT.with(|limit| {71		let current = limit.current_depth.get();72		if current < limit.max_stack_size.get() {73			limit.current_depth.set(current + 1);74			Ok(StackDepthGuard(PhantomData))75		} else {76			Err(StackOverflowError)77		}78	})79}8081pub struct StackDepthLimitOverrideGuard {82	old_limit: usize,83}84impl Drop for StackDepthLimitOverrideGuard {85	fn drop(&mut self) {86		STACK_LIMIT.with(|limit| limit.max_stack_size.set(self.old_limit));87	}88}8990pub fn limit_stack_depth(depth_limit: usize) -> StackDepthLimitOverrideGuard {91	STACK_LIMIT.with(|limit| {92		let old_limit = limit.max_stack_size.get();93		let current_depth = limit.current_depth.get();9495		limit.max_stack_size.set(current_depth + depth_limit);96		StackDepthLimitOverrideGuard { old_limit }97	})98}99100/// Like [`limit_stack_depth`], but set depth is not guarded, and will be kept101///102/// Used to implement `set_max_stack` in C api, prefer to use [`limit_stack_depth`] instead103pub fn set_stack_depth_limit(depth_limit: usize) {104	std::mem::forget(limit_stack_depth(depth_limit));105}