Skip to main content

sim_lib_agent_runner_core/
grammar.rs

1use sim_kernel::{Error, Expr, Result};
2use sim_shape::{
3    AnyShape, ExactExprShape, ExprKind, ExprKindShape, FieldShape, ListShape, OneOfShape, Shape,
4};
5
6/// Lowers a SIM `shape` into a constrained-decoding grammar string.
7///
8/// The grammar can be handed to a model runner to constrain generation to
9/// values that match `shape`.
10pub fn shape_to_grammar(shape: &dyn Shape) -> Result<String> {
11    lower_shape(shape)
12}
13
14fn lower_shape(shape: &dyn Shape) -> Result<String> {
15    if shape.as_any().is::<AnyShape>() {
16        return Ok("true".to_owned());
17    }
18    if let Some(kind) = shape.as_any().downcast_ref::<ExprKindShape>() {
19        return lower_expr_kind(kind.kind());
20    }
21    if let Some(fields) = shape.as_any().downcast_ref::<FieldShape>() {
22        return lower_field_shape(fields);
23    }
24    if let Some(list) = shape.as_any().downcast_ref::<ListShape>() {
25        return lower_list_shape(list);
26    }
27    if let Some(one_of) = shape.as_any().downcast_ref::<OneOfShape>() {
28        let choices = one_of
29            .choices()
30            .iter()
31            .map(|choice| lower_shape(choice.as_ref()))
32            .collect::<Result<Vec<_>>>()?;
33        return Ok(format!(r#"{{"anyOf":[{}]}}"#, choices.join(",")));
34    }
35    if let Some(exact) = shape.as_any().downcast_ref::<ExactExprShape>() {
36        return Ok(format!(r#"{{"const":{}}}"#, json_expr(exact.expected())?));
37    }
38    Err(Error::Eval(
39        "shape_to_grammar does not support this shape".to_owned(),
40    ))
41}
42
43fn lower_expr_kind(kind: &ExprKind) -> Result<String> {
44    Ok(match kind {
45        ExprKind::Nil => r#"{"type":"null"}"#.to_owned(),
46        ExprKind::Bool => r#"{"type":"boolean"}"#.to_owned(),
47        ExprKind::Number => r#"{"type":"number"}"#.to_owned(),
48        ExprKind::String => r#"{"type":"string"}"#.to_owned(),
49        ExprKind::List | ExprKind::Vector => r#"{"type":"array"}"#.to_owned(),
50        ExprKind::Map => r#"{"type":"object"}"#.to_owned(),
51        ExprKind::Symbol => r#"{"type":"string","description":"symbol"}"#.to_owned(),
52        other => {
53            return Err(Error::Eval(format!(
54                "shape_to_grammar does not support expr-kind {}",
55                other.name()
56            )));
57        }
58    })
59}
60
61fn lower_field_shape(shape: &FieldShape) -> Result<String> {
62    let properties = shape
63        .fields()
64        .iter()
65        .map(|field| {
66            Ok(format!(
67                "{}:{}",
68                json_string(field.name().name.as_ref()),
69                lower_shape(field.shape().as_ref())?
70            ))
71        })
72        .collect::<Result<Vec<_>>>()?;
73    let required = shape
74        .fields()
75        .iter()
76        .map(|field| json_string(field.name().name.as_ref()))
77        .collect::<Vec<_>>();
78    Ok(format!(
79        r#"{{"type":"object","properties":{{{}}},"required":[{}],"additionalProperties":false}}"#,
80        properties.join(","),
81        required.join(","),
82    ))
83}
84
85fn lower_list_shape(shape: &ListShape) -> Result<String> {
86    let prefix_items = shape
87        .items()
88        .iter()
89        .map(|item| lower_shape(item.as_ref()))
90        .collect::<Result<Vec<_>>>()?;
91    let items = match shape.rest() {
92        Some(rest) => lower_shape(rest.as_ref())?,
93        None => "false".to_owned(),
94    };
95    let bounds = if shape.rest().is_none() {
96        format!(
97            r#","minItems":{},"maxItems":{}"#,
98            shape.items().len(),
99            shape.items().len()
100        )
101    } else {
102        String::new()
103    };
104    Ok(format!(
105        r#"{{"type":"array","prefixItems":[{}],"items":{}{}}}"#,
106        prefix_items.join(","),
107        items,
108        bounds,
109    ))
110}
111
112fn json_expr(expr: &Expr) -> Result<String> {
113    Ok(match expr {
114        Expr::Nil => "null".to_owned(),
115        Expr::Bool(value) => value.to_string(),
116        Expr::Number(number) => number.canonical.clone(),
117        Expr::String(text) => json_string(text),
118        Expr::Symbol(symbol) => json_string(&symbol.to_string()),
119        Expr::List(items) | Expr::Vector(items) => {
120            let items = items.iter().map(json_expr).collect::<Result<Vec<_>>>()?;
121            format!("[{}]", items.join(","))
122        }
123        Expr::Map(entries) => {
124            let entries = entries
125                .iter()
126                .map(|(key, value)| {
127                    let Expr::Symbol(symbol) = key else {
128                        return Err(Error::Eval(
129                            "shape_to_grammar exact map keys must be symbols".to_owned(),
130                        ));
131                    };
132                    Ok(format!(
133                        "{}:{}",
134                        json_string(symbol.name.as_ref()),
135                        json_expr(value)?
136                    ))
137                })
138                .collect::<Result<Vec<_>>>()?;
139            format!("{{{}}}", entries.join(","))
140        }
141        _ => {
142            return Err(Error::Eval(
143                "shape_to_grammar exact expr lowering only supports json-like forms".to_owned(),
144            ));
145        }
146    })
147}
148
149fn json_string(text: &str) -> String {
150    format!("{text:?}")
151}
152
153#[cfg(test)]
154mod tests {
155    use super::shape_to_grammar;
156    use sim_kernel::Symbol;
157    use sim_shape::{ExprKind, ExprKindShape, FieldShape, FieldSpec, ListShape};
158    use std::sync::Arc;
159
160    #[test]
161    fn lowers_non_trivial_object_shape() {
162        let grammar = shape_to_grammar(&FieldShape::anonymous(vec![
163            FieldSpec::required(
164                Symbol::new("name"),
165                Arc::new(ExprKindShape::new(ExprKind::String)),
166            ),
167            FieldSpec::required(
168                Symbol::new("versions"),
169                Arc::new(ListShape::new(vec![
170                    Arc::new(ExprKindShape::new(ExprKind::String)),
171                    Arc::new(ExprKindShape::new(ExprKind::String)),
172                ])),
173            ),
174        ]))
175        .unwrap();
176
177        assert!(grammar.contains(r#""type":"object""#));
178        assert!(grammar.contains(r#""name":{"type":"string"}"#));
179        assert!(grammar.contains(r#""versions":{"type":"array""#));
180        assert!(grammar.contains(r#""additionalProperties":false"#));
181    }
182}