git.delta.rocks / jrsonnet / refs/commits / 0d4913596ab0

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#[thread_local]12static STACK_LIMIT: StackLimit = StackLimit {13	max_stack_size: Cell::new(200),14	current_depth: Cell::new(0),15};16#[cfg(not(feature = "nightly"))]17thread_local! {18	static STACK_LIMIT: StackLimit = StackLimit {19		max_stack_size: Cell::new(200),20		current_depth: Cell::new(0),21	};22}2324pub struct StackOverflowError;25impl From<StackOverflowError> for ErrorKind {26	fn from(_: StackOverflowError) -> Self {27		ErrorKind::StackOverflow28	}29}30impl From<StackOverflowError> for Error {31	fn from(_: StackOverflowError) -> Self {32		ErrorKind::StackOverflow.into()33	}34}3536/// Used to implement stack depth limitation37pub struct StackDepthGuard(PhantomData<()>);38impl Drop for StackDepthGuard {39	#[cfg(feature = "nightly")]40	fn drop(&mut self) {41		STACK_LIMIT42			.current_depth43			.set(STACK_LIMIT.current_depth.get() - 1)44	}45	#[cfg(not(feature = "nightly"))]46	fn drop(&mut self) {47		STACK_LIMIT.with(|limit| limit.current_depth.set(limit.current_depth.get() - 1));48	}49}5051// #[cfg(feature = "nightly")]52pub fn check_depth() -> Result<StackDepthGuard, StackOverflowError> {53	fn internal(limit: &StackLimit) -> Result<StackDepthGuard, StackOverflowError> {54		let current = limit.current_depth.get();55		if current < limit.max_stack_size.get() {56			limit.current_depth.set(current + 1);57			Ok(StackDepthGuard(PhantomData))58		} else {59			Err(StackOverflowError)60		}61	}62	#[cfg(feature = "nightly")]63	{64		internal(&STACK_LIMIT)65	}66	#[cfg(not(feature = "nightly"))]67	{68		STACK_LIMIT.with(internal)69	}70}7172pub struct StackDepthLimitOverrideGuard {73	old_limit: usize,74}75impl Drop for StackDepthLimitOverrideGuard {76	#[cfg(feature = "nightly")]77	fn drop(&mut self) {78		STACK_LIMIT.max_stack_size.set(self.old_limit)79	}80	#[cfg(not(feature = "nightly"))]81	fn drop(&mut self) {82		STACK_LIMIT.with(|limit| limit.max_stack_size.set(self.old_limit));83	}84}8586pub fn limit_stack_depth(depth_limit: usize) -> StackDepthLimitOverrideGuard {87	fn internal(limit: &StackLimit, depth_limit: usize) -> StackDepthLimitOverrideGuard {88		let old_limit = limit.max_stack_size.get();89		let current_depth = limit.current_depth.get();9091		limit.max_stack_size.set(current_depth + depth_limit);92		StackDepthLimitOverrideGuard { old_limit }93	}94	#[cfg(feature = "nightly")]95	{96		internal(&STACK_LIMIT, depth_limit)97	}98	#[cfg(not(feature = "nightly"))]99	{100		STACK_LIMIT.with(|limit| internal(limit, depth_limit))101	}102}103104/// Like [`limit_stack_depth`], but set depth is not guarded, and will be kept105///106/// Used to implement `set_max_stack` in C api, prefer to use [`limit_stack_depth`] instead107pub fn set_stack_depth_limit(depth_limit: usize) {108	std::mem::forget(limit_stack_depth(depth_limit));109}