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

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(feature = "nightly")]11#[allow(clippy::thread_local_initializer_can_be_made_const)]12#[thread_local]13static STACK_LIMIT: StackLimit = StackLimit {14	max_stack_size: Cell::new(200),15	current_depth: Cell::new(0),16};17#[cfg(not(feature = "nightly"))]18thread_local! {19	static STACK_LIMIT: StackLimit = const {20		StackLimit {21			max_stack_size: Cell::new(200),22			current_depth: Cell::new(0),23		}24	};25}2627pub struct StackOverflowError;28impl From<StackOverflowError> for ErrorKind {29	fn from(_: StackOverflowError) -> Self {30		Self::StackOverflow31	}32}33impl From<StackOverflowError> for Error {34	fn from(_: StackOverflowError) -> Self {35		ErrorKind::StackOverflow.into()36	}37}3839/// Used to implement stack depth limitation40pub struct StackDepthGuard(PhantomData<()>);41impl Drop for StackDepthGuard {42	#[cfg(feature = "nightly")]43	fn drop(&mut self) {44		STACK_LIMIT45			.current_depth46			.set(STACK_LIMIT.current_depth.get() - 1);47	}48	#[cfg(not(feature = "nightly"))]49	fn drop(&mut self) {50		STACK_LIMIT.with(|limit| limit.current_depth.set(limit.current_depth.get() - 1));51	}52}5354// #[cfg(feature = "nightly")]55pub fn check_depth() -> Result<StackDepthGuard, StackOverflowError> {56	fn internal(limit: &StackLimit) -> Result<StackDepthGuard, StackOverflowError> {57		let current = limit.current_depth.get();58		if current < limit.max_stack_size.get() {59			limit.current_depth.set(current + 1);60			Ok(StackDepthGuard(PhantomData))61		} else {62			Err(StackOverflowError)63		}64	}65	#[cfg(feature = "nightly")]66	{67		internal(&STACK_LIMIT)68	}69	#[cfg(not(feature = "nightly"))]70	{71		STACK_LIMIT.with(internal)72	}73}7475pub struct StackDepthLimitOverrideGuard {76	old_limit: usize,77}78impl Drop for StackDepthLimitOverrideGuard {79	#[cfg(feature = "nightly")]80	fn drop(&mut self) {81		STACK_LIMIT.max_stack_size.set(self.old_limit);82	}83	#[cfg(not(feature = "nightly"))]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	fn internal(limit: &StackLimit, depth_limit: usize) -> StackDepthLimitOverrideGuard {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	#[cfg(feature = "nightly")]98	{99		internal(&STACK_LIMIT, depth_limit)100	}101	#[cfg(not(feature = "nightly"))]102	{103		STACK_LIMIT.with(|limit| internal(limit, depth_limit))104	}105}106107/// Like [`limit_stack_depth`], but set depth is not guarded, and will be kept108///109/// Used to implement `set_max_stack` in C api, prefer to use [`limit_stack_depth`] instead110pub fn set_stack_depth_limit(depth_limit: usize) {111	std::mem::forget(limit_stack_depth(depth_limit));112}