1#![cfg_attr(not(feature = "std"), no_std)]
17
18use frame_support::dispatch::{DispatchInfo, PostDispatchInfo};
19use frame_support::traits::Get;
20use frame_system::limits::BlockWeights;
21use frame_system::{Config, ConsumedWeight};
22use parity_scale_codec::{Decode, Encode};
23use scale_info::TypeInfo;
24use sp_runtime::traits::{
25 DispatchInfoOf, DispatchOriginOf, Dispatchable, PostDispatchInfoOf, TransactionExtension,
26 ValidateResult,
27};
28use sp_runtime::transaction_validity::{
29 TransactionSource, TransactionValidity, TransactionValidityError,
30};
31use sp_runtime::DispatchResult;
32use sp_weights::Weight;
33
34#[derive(Encode, Decode, Clone, Eq, PartialEq, Default, TypeInfo)]
40#[scale_info(skip_type_params(T))]
41pub struct CheckWeight<T: Config + Send + Sync>(core::marker::PhantomData<T>);
42
43impl<T: Config + Send + Sync> CheckWeight<T>
44where
45 T::RuntimeCall: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
46{
47 pub fn new() -> Self {
49 Self(Default::default())
50 }
51
52 pub fn do_prepare(
57 info: &DispatchInfoOf<T::RuntimeCall>,
58 len: usize,
59 next_len: u32,
60 ) -> Result<(), TransactionValidityError> {
61 let all_weight = frame_system::Pallet::<T>::block_weight();
62 let maximum_weight = T::BlockWeights::get();
63 let next_weight =
64 calculate_consumed_weight::<T::RuntimeCall>(&maximum_weight, all_weight, info, len);
65
66 frame_system::AllExtrinsicsLen::<T>::put(next_len);
67 frame_system::BlockWeight::<T>::put(next_weight);
68
69 Ok(())
70 }
71}
72
73fn calculate_consumed_weight<Call>(
75 maximum_weight: &BlockWeights,
76 mut all_weight: ConsumedWeight,
77 info: &DispatchInfoOf<Call>,
78 len: usize,
79) -> ConsumedWeight
80where
81 Call: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
82{
83 let extrinsic_weight = (info.call_weight + info.extension_weight)
85 .saturating_add(maximum_weight.get(info.class).base_extrinsic)
86 .saturating_add(Weight::from_parts(0, len as u64));
87
88 all_weight.accrue(extrinsic_weight, info.class);
90
91 all_weight
92}
93
94impl<T: Config + Send + Sync> TransactionExtension<T::RuntimeCall> for CheckWeight<T>
95where
96 T::RuntimeCall: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
97{
98 type Implicit = ();
99 type Pre = ();
100 type Val = u32;
101 const IDENTIFIER: &'static str = "CheckWeight";
102
103 fn weight(&self, _: &T::RuntimeCall) -> Weight {
104 <T::ExtensionsWeightInfo as frame_system::ExtensionsWeightInfo>::check_weight()
105 }
106
107 fn validate(
108 &self,
109 origin: T::RuntimeOrigin,
110 _call: &T::RuntimeCall,
111 info: &DispatchInfoOf<T::RuntimeCall>,
112 len: usize,
113 _self_implicit: Self::Implicit,
114 _inherited_implication: &impl Encode,
115 _source: TransactionSource,
116 ) -> ValidateResult<Self::Val, T::RuntimeCall> {
117 let (validity, next_len) = frame_system::CheckWeight::<T>::do_validate(info, len)?;
118 Ok((validity, next_len, origin))
119 }
120
121 fn prepare(
122 self,
123 val: Self::Val,
124 _origin: &DispatchOriginOf<T::RuntimeCall>,
125 _call: &T::RuntimeCall,
126 info: &DispatchInfoOf<T::RuntimeCall>,
127 len: usize,
128 ) -> Result<Self::Pre, TransactionValidityError> {
129 Self::do_prepare(info, len, val)
130 }
131
132 fn post_dispatch_details(
133 _pre: Self::Pre,
134 info: &DispatchInfoOf<T::RuntimeCall>,
135 post_info: &PostDispatchInfoOf<T::RuntimeCall>,
136 _len: usize,
137 _result: &DispatchResult,
138 ) -> Result<Weight, TransactionValidityError> {
139 frame_system::CheckWeight::<T>::do_post_dispatch(info, post_info)?;
140 Ok(Weight::zero())
141 }
142
143 fn bare_validate(
144 _call: &T::RuntimeCall,
145 info: &DispatchInfoOf<T::RuntimeCall>,
146 len: usize,
147 ) -> TransactionValidity {
148 Ok(frame_system::CheckWeight::<T>::do_validate(info, len)?.0)
149 }
150
151 fn bare_validate_and_prepare(
152 _call: &T::RuntimeCall,
153 info: &DispatchInfoOf<T::RuntimeCall>,
154 len: usize,
155 ) -> Result<(), TransactionValidityError> {
156 let (_, next_len) = frame_system::CheckWeight::<T>::do_validate(info, len)?;
157 Self::do_prepare(info, len, next_len)
158 }
159
160 fn bare_post_dispatch(
161 info: &DispatchInfoOf<T::RuntimeCall>,
162 post_info: &mut PostDispatchInfoOf<T::RuntimeCall>,
163 _len: usize,
164 _result: &DispatchResult,
165 ) -> Result<(), TransactionValidityError> {
166 frame_system::CheckWeight::<T>::do_post_dispatch(info, post_info)
167 }
168}
169
170impl<T: Config + Send + Sync> core::fmt::Debug for CheckWeight<T> {
171 #[cfg(feature = "std")]
172 fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
173 write!(f, "CheckWeight")
174 }
175
176 #[cfg(not(feature = "std"))]
177 fn fmt(&self, _: &mut core::fmt::Formatter) -> core::fmt::Result {
178 Ok(())
179 }
180}