1use std::fmt::Debug;
2
3use trellis_core::{
4 AuditExplanationLevel, Graph, GraphResult, InvariantResultTrace, OutputFrameTrace,
5 ResourceCommandTrace, Transaction, TransactionOptions,
6};
7
8use crate::harness_step::{HarnessStep, NamedInvariantCheck};
9use crate::{
10 DataTransactionScript, FullRecomputeOracle, OracleCheck, OracleMismatch, OutputLedger,
11 ResourceLedger, Scenario, ScenarioError, StageOperation, TransactionScript,
12};
13
14pub trait ScenarioTarget<C = ()> {
16 fn graph(&self) -> &Graph<C>;
18
19 fn graph_mut(&mut self) -> &mut Graph<C>;
21}
22
23impl<C> ScenarioTarget<C> for Graph<C> {
24 fn graph(&self) -> &Graph<C> {
25 self
26 }
27
28 fn graph_mut(&mut self) -> &mut Graph<C> {
29 self
30 }
31}
32
33pub struct TrellisHarness<G, C = ()> {
35 target: G,
36 scenario: Scenario,
37 resource_ledger: ResourceLedger<C>,
38 output_ledger: OutputLedger,
39}
40
41impl<G, C> TrellisHarness<G, C>
42where
43 G: ScenarioTarget<C>,
44 C: Clone + Debug + PartialEq,
45{
46 pub fn new(build: impl FnOnce() -> G) -> Self {
48 Self::from_target(build())
49 }
50
51 pub fn from_target(target: G) -> Self {
53 Self {
54 target,
55 scenario: Scenario::new(),
56 resource_ledger: ResourceLedger::new(),
57 output_ledger: OutputLedger::new(),
58 }
59 }
60
61 pub fn target(&self) -> &G {
63 &self.target
64 }
65
66 pub fn scenario(&self) -> &Scenario {
68 &self.scenario
69 }
70
71 pub fn resource_ledger(&self) -> &ResourceLedger<C> {
73 &self.resource_ledger
74 }
75
76 pub fn output_ledger(&self) -> &OutputLedger {
78 &self.output_ledger
79 }
80
81 pub fn step(&mut self, name: impl Into<String>) -> HarnessStep<'_, G, C> {
83 HarnessStep::new(self, name.into())
84 }
85
86 pub fn run_script(&mut self, script: &TransactionScript<C>) -> Result<(), ScenarioError> {
88 for step in script.steps() {
89 self.commit_operations(step.name(), &step.operations, &[], None, None)?;
90 }
91 Ok(())
92 }
93
94 pub fn run_data_script<Operation>(
96 &mut self,
97 script: &DataTransactionScript<Operation>,
98 mut apply: impl for<'tx> FnMut(&Operation, &mut Transaction<'tx, C>) -> GraphResult<()>,
99 ) -> Result<(), ScenarioError> {
100 script.validate_format_version()?;
101 for step in script.steps() {
102 self.commit_data_operations(step.name(), step.operations(), &mut apply)?;
103 }
104 Ok(())
105 }
106
107 pub fn replay(
109 build: impl FnOnce() -> G,
110 script: &TransactionScript<C>,
111 ) -> Result<Self, ScenarioError> {
112 let mut harness = Self::new(build);
113 harness.run_script(script)?;
114 Ok(harness)
115 }
116
117 pub fn replay_data<Operation>(
119 build: impl FnOnce() -> G,
120 script: &DataTransactionScript<Operation>,
121 apply: impl for<'tx> FnMut(&Operation, &mut Transaction<'tx, C>) -> GraphResult<()>,
122 ) -> Result<Self, ScenarioError> {
123 let mut harness = Self::new(build);
124 harness.run_data_script(script, apply)?;
125 Ok(harness)
126 }
127
128 pub fn assert_replay_matches(&self, other: &Self) -> Result<(), ScenarioError> {
130 self.scenario.assert_replay_matches(&other.scenario)?;
131 let expected = self.final_state_debug_dump();
132 let actual = other.final_state_debug_dump();
133 if expected != actual {
134 return Err(ScenarioError::ReplayFinalStateMismatch { expected, actual });
135 }
136 assert_equal_debug(
137 "resource_command_records",
138 self.resource_ledger.command_records(),
139 other.resource_ledger.command_records(),
140 )?;
141 assert_equal_debug(
142 "output_frame_records",
143 self.output_ledger.frame_records(),
144 other.output_ledger.frame_records(),
145 )?;
146 assert_equal_debug(
147 "resource_ledger_snapshots",
148 &self.resource_ledger,
149 &other.resource_ledger,
150 )?;
151 assert_equal_debug(
152 "output_ledger_snapshots",
153 &self.output_ledger,
154 &other.output_ledger,
155 )?;
156 Ok(())
157 }
158
159 pub fn final_state_debug_dump(&self) -> String {
161 self.target.graph().debug_dump()
162 }
163
164 pub fn assert_oracle<Oracle>(
166 &self,
167 inputs: &Oracle::CanonicalInputs,
168 ) -> Result<OracleCheck<Oracle::ExpectedState>, OracleMismatch<Oracle::ExpectedState>>
169 where
170 Oracle: FullRecomputeOracle<G>,
171 {
172 crate::assert_incremental_equals_full::<G, Oracle>(&self.target, inputs)
173 }
174
175 pub(crate) fn commit_operations(
176 &mut self,
177 name: &str,
178 operations: &[Box<StageOperation<C>>],
179 invariant_checks: &[NamedInvariantCheck<G, C>],
180 expected_resource_commands: Option<&[ResourceCommandTrace]>,
181 expected_output_frames: Option<&[OutputFrameTrace]>,
182 ) -> Result<(), ScenarioError> {
183 self.scenario.ensure_step_name_available(name)?;
184 let result = {
185 let graph = self.target.graph_mut();
186 let mut tx = graph
187 .begin_transaction_with_options(harness_transaction_options())
188 .map_err(|error| step_commit_failed(name, error))?;
189 for operation in operations {
190 operation(&mut tx).map_err(|error| step_commit_failed(name, error))?;
191 }
192 tx.commit()
193 .map_err(|error| step_commit_failed(name, error))?
194 };
195
196 let mut trace = result.trace();
197 for check in invariant_checks {
198 let passed = (check.check)(&self.target, &result);
199 trace.invariant_results.push(InvariantResultTrace {
200 name: check.name.clone(),
201 passed,
202 });
203 if !passed {
204 return Err(ScenarioError::InvariantFailed {
205 step: name.to_owned(),
206 invariant: check.name.clone(),
207 transaction_id: result.transaction_id,
208 revision: result.revision,
209 });
210 }
211 }
212
213 self.resource_ledger.apply_result(&result);
214 self.output_ledger.apply_result(&result);
215 self.resource_ledger
216 .assert_graph_has_no_orphan_resources(self.target.graph())
217 .map_err(|error| ScenarioError::ResourceLedgerInvariantFailed {
218 step: name.to_owned(),
219 error: Box::new(error),
220 })?;
221 self.scenario.record_trace(name, trace)?;
222
223 if let Some(expected) = expected_resource_commands {
224 self.scenario
225 .assert_step_resource_commands(name, expected)?;
226 }
227 if let Some(expected) = expected_output_frames {
228 self.scenario.assert_step_output_frames(name, expected)?;
229 }
230 Ok(())
231 }
232
233 fn commit_data_operations<Operation>(
234 &mut self,
235 name: &str,
236 operations: &[Operation],
237 apply: &mut impl for<'tx> FnMut(&Operation, &mut Transaction<'tx, C>) -> GraphResult<()>,
238 ) -> Result<(), ScenarioError> {
239 self.scenario.ensure_step_name_available(name)?;
240 let result = {
241 let graph = self.target.graph_mut();
242 let mut tx = graph
243 .begin_transaction_with_options(harness_transaction_options())
244 .map_err(|error| step_commit_failed(name, error))?;
245 for operation in operations {
246 apply(operation, &mut tx).map_err(|error| step_commit_failed(name, error))?;
247 }
248 tx.commit()
249 .map_err(|error| step_commit_failed(name, error))?
250 };
251
252 self.resource_ledger.apply_result(&result);
253 self.output_ledger.apply_result(&result);
254 self.resource_ledger
255 .assert_graph_has_no_orphan_resources(self.target.graph())
256 .map_err(|error| ScenarioError::ResourceLedgerInvariantFailed {
257 step: name.to_owned(),
258 error: Box::new(error),
259 })?;
260 self.scenario.record(name, &result)
261 }
262}
263
264fn harness_transaction_options() -> TransactionOptions {
265 TransactionOptions::default().with_audit_explanations(AuditExplanationLevel::DependencyPaths)
266}
267
268fn step_commit_failed(step: &str, error: trellis_core::GraphError) -> ScenarioError {
269 ScenarioError::StepCommitFailed {
270 step: step.to_owned(),
271 error,
272 }
273}
274
275fn assert_equal_debug<T>(field: &'static str, expected: &T, actual: &T) -> Result<(), ScenarioError>
276where
277 T: Debug + PartialEq + ?Sized,
278{
279 if expected == actual {
280 Ok(())
281 } else {
282 Err(ScenarioError::ReplayLedgerMismatch {
283 field,
284 expected: format!("{expected:#?}"),
285 actual: format!("{actual:#?}"),
286 })
287 }
288}