GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/openvmm

Public

mirrored fromhttps://github.com/microsoft/openvmmAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
e6a1eca0ac970ac1b75f1b4ea3ba48c126ddd5ea

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

openvmm/flowey/flowey_core/src

flowey/flowey_core/src/patch.rs

228lines · modecode

RawDownload
Latest commit unavailable.
1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4use crate::node::FlowNodeBase;
5use crate::node::NodeHandle;
6use crate::node::WriteVar;
7use std::collections::BTreeMap;
8use std::sync::OnceLock;
9
10pub type PatchFn = fn(&mut PatchManager<'_>);
11
12// A patchfn that does nothing. Can be useful when writing logic that
13// conditionally applies patches.
14pub fn noop_patchfn(_: &mut PatchManager<'_>) {}
15
16enum PatchEvent {
17 Swap {
18 from_old_node: NodeHandle,
19 with_new_node: NodeHandle,
20 },
21 InjectSideEffect {
22 from_old_node: NodeHandle,
23 with_new_node: NodeHandle,
24 side_effect_var: String,
25 req: Box<[u8]>,
26 },
27}
28
29trait PatchManagerBackend {
30 fn new_side_effect_var(&mut self) -> String;
31 fn on_patch_event(&mut self, event: PatchEvent);
32}
33
34/// Passed to patch functions
35pub struct PatchManager<'a> {
36 backend: &'a mut dyn PatchManagerBackend,
37}
38
39impl PatchManager<'_> {
40 pub fn hook<N: FlowNodeBase>(&mut self) -> PatchHook<'_, N> {
41 PatchHook {
42 backend: self.backend,
43 _kind: std::marker::PhantomData,
44 }
45 }
46}
47
48/// Patch operations in the context of a particular Node.
49pub struct PatchHook<'a, N: FlowNodeBase> {
50 backend: &'a mut dyn PatchManagerBackend,
51 _kind: std::marker::PhantomData<N>,
52}
53
54impl<N> PatchHook<'_, N>
55where
56 N: FlowNodeBase + 'static,
57{
58 /// Swap out the target Node's implementation with a different
59 /// implementation.
60 pub fn swap_with<M>(&mut self) -> &mut Self
61 where
62 M: 'static,
63 // use the type system to enforce that patch nodes have an identical
64 // request type
65 M: FlowNodeBase<Request = N::Request>,
66 {
67 self.backend.on_patch_event(PatchEvent::Swap {
68 from_old_node: NodeHandle::from_type::<N>(),
69 with_new_node: NodeHandle::from_type::<M>(),
70 });
71 self
72 }
73
74 /// Inject a side-effect dependency, which runs before any other steps in
75 /// the Node.
76 pub fn inject_side_effect<T, M>(
77 &mut self,
78 f: impl FnOnce(WriteVar<T>) -> M::Request,
79 ) -> &mut Self
80 where
81 T: serde::Serialize + serde::de::DeserializeOwned,
82 M: 'static,
83 M: FlowNodeBase,
84 {
85 let backing_var = self.backend.new_side_effect_var();
86 let req = f(crate::node::thin_air_write_runtime_var(
87 backing_var.clone(),
88 false,
89 ));
90
91 self.backend.on_patch_event(PatchEvent::InjectSideEffect {
92 from_old_node: NodeHandle::from_type::<N>(),
93 with_new_node: NodeHandle::from_type::<M>(),
94 side_effect_var: backing_var,
95 req: serde_json::to_vec(&req).map(Into::into).unwrap(),
96 });
97 self
98 }
99}
100
101pub fn patchfn_by_modpath() -> &'static BTreeMap<String, PatchFn> {
102 static MODPATH_LOOKUP: OnceLock<BTreeMap<String, PatchFn>> = OnceLock::new();
103
104 let lookup = MODPATH_LOOKUP.get_or_init(|| {
105 let mut lookup = BTreeMap::new();
106 for (f, module_path, fn_name) in private::PATCH_FNS {
107 let existing = lookup.insert(format!("{}::{}", module_path, fn_name), *f);
108 // Rust would've errored out at module defn time with a duplicate fn name error
109 assert!(existing.is_none());
110 }
111 lookup
112 });
113
114 lookup
115}
116
117/// [`PatchResolver`]
118#[derive(Debug, Clone)]
119pub struct ResolvedPatches {
120 pub swap: BTreeMap<NodeHandle, NodeHandle>,
121 pub inject_side_effect: BTreeMap<NodeHandle, Vec<(NodeHandle, String, Box<[u8]>)>>,
122}
123
124impl ResolvedPatches {
125 pub fn build() -> PatchResolver {
126 PatchResolver {
127 side_effect_var_idx: 0,
128 swap: BTreeMap::default(),
129 inject_side_effect: BTreeMap::new(),
130 }
131 }
132}
133
134/// Helper method to resolve multiple patches into a single [`ResolvedPatches`]
135#[derive(Debug)]
136pub struct PatchResolver {
137 side_effect_var_idx: usize,
138 swap: BTreeMap<NodeHandle, NodeHandle>,
139 inject_side_effect: BTreeMap<NodeHandle, Vec<(NodeHandle, String, Box<[u8]>)>>,
140}
141
142impl PatchResolver {
143 pub fn apply_patchfn(&mut self, patchfn: PatchFn) {
144 patchfn(&mut PatchManager { backend: self });
145 }
146
147 pub fn finalize(self) -> ResolvedPatches {
148 let Self {
149 swap,
150 mut inject_side_effect,
151 side_effect_var_idx: _,
152 } = self;
153
154 // take into account the interaction between swaps and injected effects
155 for (from, to) in &swap {
156 let injected = inject_side_effect.remove(from);
157 if let Some(injected) = injected {
158 inject_side_effect.insert(*to, injected);
159 }
160 }
161
162 ResolvedPatches {
163 swap,
164 inject_side_effect,
165 }
166 }
167}
168
169impl PatchManagerBackend for PatchResolver {
170 fn new_side_effect_var(&mut self) -> String {
171 self.side_effect_var_idx += 1;
172 format!("patch_side_effect:{}", self.side_effect_var_idx)
173 }
174
175 fn on_patch_event(&mut self, event: PatchEvent) {
176 match event {
177 PatchEvent::Swap {
178 from_old_node,
179 with_new_node,
180 } => {
181 let existing = self.swap.insert(from_old_node, with_new_node);
182 // FUTURE: add some better error reporting / logging to
183 // allow doing this, albeit with a warning
184 assert!(
185 existing.is_none(),
186 "cannot double-patch the same node combo"
187 );
188 }
189 PatchEvent::InjectSideEffect {
190 from_old_node,
191 with_new_node,
192 side_effect_var,
193 req,
194 } => {
195 self.inject_side_effect
196 .entry(from_old_node)
197 .or_default()
198 .push((with_new_node, side_effect_var, req));
199 }
200 }
201 }
202}
203
204#[doc(hidden)]
205pub mod private {
206 use super::PatchFn;
207 pub use linkme;
208
209 #[linkme::distributed_slice]
210 pub static PATCH_FNS: [(PatchFn, &'static str, &'static str)] = [..];
211
212 /// Register a patch function which can be used when emitting flows.
213 ///
214 /// The function must conform to the signature of [`PatchFn`]
215 #[macro_export]
216 macro_rules! register_patch {
217 ($patchfn:ident) => {
218 const _: () = {
219 use $crate::node::private::linkme;
220
221 #[linkme::distributed_slice($crate::patch::private::PATCH_FNS)]
222 #[linkme(crate = linkme)]
223 pub static PATCH_FNS: ($crate::patch::PatchFn, &'static str, &'static str) =
224 ($patchfn, module_path!(), stringify!($patchfn));
225 };
226 };
227 }
228}