1use serde_json::{Map, Value, json};
4
5use henad_compute::entry::ModelEntry;
6use henad_core::explore::fingerprint::schema_hash;
7use henad_core::metadata::Backend;
8use henad_core::params::{ParamApply, ParamDescriptor, ParamFormat, ParamKind};
9
10use crate::probe::ProbeReport;
11
12pub const SCHEMA_VERSION: u64 = 1;
14
15pub fn schema_json(entry: &ModelEntry, probe: Option<&ProbeReport>) -> Value {
20 let params: Vec<Value> = entry
21 .param_descriptors()
22 .iter()
23 .enumerate()
24 .map(|(index, descriptor)| param_json(index, descriptor))
25 .collect();
26 let stats: Vec<Value> = entry
27 .stat_descriptors()
28 .iter()
29 .map(|stat| json!({ "label": stat.label, "color": stat.color }))
30 .collect();
31 let actions: Vec<Value> = entry
32 .action_descriptors()
33 .iter()
34 .enumerate()
35 .map(|(index, action)| json!({ "index": index, "id": action.id, "label": action.label }))
36 .collect();
37 let mut schema = json!({
38 "schema_version": SCHEMA_VERSION,
39 "model": entry.id(),
40 "name": entry.name(),
41 "backend": backend_name(entry.metadata().backend),
42 "schema_hash": format!("{:016x}", schema_hash(&entry.schema())),
43 "params": params,
44 "stats": stats,
45 "actions": actions,
46 "seed": { "type": "u64" },
47 });
48 if let Some(probe) = probe {
49 let columns: Vec<&str> = (0..probe.columns.len()).map(|i| probe.columns.name(i)).collect();
50 schema["stat_columns"] = json!(columns);
51 }
52 schema
53}
54
55pub fn backend_name(backend: Backend) -> &'static str {
57 match backend {
58 Backend::Cpu => "cpu",
59 Backend::Gpu => "gpu",
60 }
61}
62
63fn param_json(index: usize, descriptor: &ParamDescriptor) -> Value {
64 let mut param = Map::new();
65 param.insert("index".to_owned(), json!(index));
66 param.insert("id".to_owned(), json!(descriptor.id));
67 param.insert("label".to_owned(), json!(descriptor.label));
68 let apply = match descriptor.apply {
69 ParamApply::Live => "live",
70 ParamApply::OnReload => "reload",
71 };
72 param.insert("apply".to_owned(), json!(apply));
73 let format = match descriptor.format {
74 ParamFormat::Plain => "plain",
75 ParamFormat::Percent => "percent",
76 };
77 param.insert("format".to_owned(), json!(format));
78 let fields = match descriptor.kind {
79 ParamKind::F32 {
80 min,
81 max,
82 default,
83 step,
84 } => {
85 let mut fields = json!({
86 "kind": "f32",
87 "min": f32_json(min),
88 "max": f32_json(max),
89 "default": f32_json(default),
90 });
91 if let Some(step) = step {
92 fields["step"] = f32_json(step);
93 }
94 fields
95 }
96 ParamKind::U32 { min, max, default } => {
97 json!({ "kind": "u32", "min": min, "max": max, "default": default })
98 }
99 ParamKind::Bool { default } => json!({ "kind": "bool", "default": default }),
100 ParamKind::Choice { options, default } => json!({
101 "kind": "choice",
102 "options": options,
103 "default": options.get(default),
104 "default_index": default,
105 }),
106 };
107 if let Value::Object(fields) = fields {
108 param.extend(fields);
109 }
110 Value::Object(param)
111}
112
113pub(crate) fn f32_json(value: f32) -> Value {
117 value
118 .to_string()
119 .parse::<f64>()
120 .ok()
121 .and_then(serde_json::Number::from_f64)
122 .map_or(Value::Null, Value::Number)
123}
124
125#[cfg(test)]
126mod tests {
127 use henad_models::example_models;
128 use serde_json::json;
129
130 use henad_core::explore::fingerprint::schema_hash;
131
132 use super::{f32_json, schema_json};
133 use crate::probe::ProbeReport;
134
135 const SCHEMA_HASHES_0_2_0: [(&str, &str); 10] = [
139 ("sir", "6ff1dc3971fd0a96"),
140 ("boids", "ae99f3e0d37c8d3d"),
141 ("game_of_life", "aba7303a467b06bf"),
142 ("ants", "f6eb31e76efdf4cf"),
143 ("virus_network", "77993c7b4047b8bd"),
144 ("team_assembly", "d1724ce65dada5d6"),
145 ("gpu_game_of_life", "461ad8e0d063a398"),
146 ("gpu_sir", "99f14d30367e3743"),
147 ("gpu_boids", "f7328729019009c0"),
148 ("gpu_ants", "b58e8a5a5b6a829e"),
149 ];
150
151 #[test]
155 fn schema_hashes_are_unchanged_since_0_2_0() {
156 let models = example_models();
157 for (id, recorded) in SCHEMA_HASHES_0_2_0 {
158 let entry = models.get(id).unwrap_or_else(|| panic!("{id} is registered"));
159 let current = format!("{:016x}", schema_hash(&entry.schema()));
160 assert_eq!(current, recorded, "{id}'s schema hash moved since 0.2.0");
161 }
162 }
163
164 #[test]
165 fn every_descriptor_is_listed_with_its_kind_and_bounds() {
166 let models = example_models();
167 for entry in models.iter().filter(|entry| entry.gpu_needs().is_none()) {
168 let schema = schema_json(entry, None);
169 let params = schema["params"].as_array().expect("params is a list");
170 assert_eq!(params.len(), entry.param_descriptors().len(), "{}", entry.id());
171 for (index, param) in params.iter().enumerate() {
172 assert_eq!(param["index"], json!(index));
173 assert_eq!(param["id"], json!(entry.param_descriptors()[index].id));
174 assert!(param["kind"].is_string() && param["default"] != json!(null), "{param}");
175 }
176 assert_eq!(schema["model"], json!(entry.id()));
177 assert_eq!(schema["schema_hash"].as_str().map(str::len), Some(16));
178 assert!(schema.get("stat_columns").is_none(), "no probe, no columns");
179 }
180 }
181
182 #[test]
183 fn a_probe_adds_the_stat_columns() {
184 let entry = henad_models::example_models()
185 .get("sir")
186 .cloned()
187 .expect("sir is registered");
188 let defaults: Vec<_> = entry
189 .param_descriptors()
190 .iter()
191 .map(|descriptor| descriptor.kind.default_value())
192 .collect();
193 let probe = ProbeReport::build(&entry, None, &defaults, None).expect("sir builds at its defaults");
194 let schema = schema_json(&entry, Some(&probe));
195 assert_eq!(schema["stat_columns"], json!(["Susceptible", "Infected", "Recovered"]));
196 let grid_width = &schema["params"][0];
197 assert_eq!(grid_width["id"], json!("grid_width"));
198 assert_eq!(grid_width["kind"], json!("u32"));
199 assert_eq!(grid_width["apply"], json!("reload"));
200 }
201
202 #[test]
203 fn an_f32_is_written_in_its_shortest_form() {
204 assert_eq!(f32_json(0.025).to_string(), "0.025");
205 assert_eq!(f32_json(1.0).to_string(), "1.0");
206 assert_eq!(f32_json(f32::NAN), json!(null));
207 }
208}