hax_rust_engine/
ocaml_engine.rs

1//! This module implements an interface to the OCaml hax engine. Via this
2//! interface, the rust engine can communicate with the OCaml engine, and reuse
3//! some of its components.
4
5use std::io::BufRead;
6
7use hax_frontend_exporter::{
8    ThirBody,
9    id_table::{Table, WithTable},
10};
11use hax_types::engine_api::{
12    EngineOptions,
13    protocol::{FromEngine, ToEngine},
14};
15use serde::Deserialize;
16
17/// A query for the OCaml engine
18#[derive(Debug, Clone, ::schemars::JsonSchema, ::serde::Deserialize, ::serde::Serialize)]
19pub struct Query {
20    /// The version of hax currently used
21    pub hax_version: String,
22    /// Dictionary from `DefId`s to `impl_infos`
23    pub impl_infos: Vec<(
24        hax_frontend_exporter::DefId,
25        hax_frontend_exporter::ImplInfos,
26    )>,
27    /// The kind of query we want to send to the engine
28    pub kind: QueryKind,
29}
30
31/// The payload of the query. [`Response`] below mirrors this enum to represent
32/// the response from the engine.
33#[derive(Debug, Clone, ::schemars::JsonSchema, ::serde::Deserialize, ::serde::Serialize)]
34pub enum QueryKind {
35    /// Ask the OCaml engine to import the given THIR from the frontend
36    ImportThir {
37        /// The input THIR items
38        input: Vec<hax_frontend_exporter::Item<ThirBody>>,
39        /// Temporary option to enable a set of default phases
40        apply_phases: bool,
41        /// Translation options which contains include clauses (items filtering)
42        translation_options: hax_types::cli_options::TranslationOptions,
43    },
44}
45
46/// A Response after a [`Query`]
47#[derive(Debug, Clone, ::schemars::JsonSchema, ::serde::Deserialize, ::serde::Serialize)]
48pub enum Response {
49    /// Return imported THIR as an internal AST from Rust engine
50    ImportThir {
51        /// The output Rust AST items
52        output: Vec<crate::ast::Item>,
53    },
54}
55
56/// Extends the common `ToEngine` messages with one extra case: `Query`.
57#[derive(::serde::Deserialize, ::serde::Serialize)]
58#[serde(untagged)]
59pub enum ExtendedToEngine {
60    /// A standard `ToEngine` message
61    ToEngine(ToEngine),
62    /// A `Query`
63    Query(Box<WithTable<EngineOptions>>),
64}
65
66/// Extends the common `FromEngine` messages with one extra case: `Response`.
67#[derive(Debug, Clone, ::schemars::JsonSchema, ::serde::Deserialize, ::serde::Serialize)]
68#[serde(untagged)]
69pub enum ExtendedFromEngine {
70    /// A standard `FromEngine` message
71    FromEngine(FromEngine),
72    /// A `Response`
73    Response(Response),
74}
75
76impl Query {
77    /// Execute the query synchronously.
78    pub fn execute(&self, table: Table) -> Option<Response> {
79        use std::io::Write;
80        use std::process::Command;
81
82        macro_rules! send {
83            ($where: expr, $value:expr) => {
84                serde_json::to_writer(&mut $where, $value).unwrap();
85                $where.write_all(b"\n").unwrap();
86                $where.flush().unwrap();
87            };
88        }
89
90        let mut engine_subprocess =
91            Command::new(std::env::var("HAX_ENGINE_BINARY").unwrap_or("hax-engine".into()))
92                .arg("driver_rust_engine")
93                .stdin(std::process::Stdio::piped())
94                .stdout(std::process::Stdio::piped())
95                .spawn()
96                .unwrap();
97
98        let mut stdin = std::io::BufWriter::new(
99            engine_subprocess
100                .stdin
101                .as_mut()
102                .expect("Could not write on stdin"),
103        );
104
105        WithTable::run(table, self, |with_table| {
106            send!(stdin, with_table);
107        });
108
109        let mut response = None;
110        let stdout = std::io::BufReader::new(engine_subprocess.stdout.take().unwrap());
111        // TODO: this should be streaming (i.e. use a `LineAsEOF` reader wrapper that consumes a reader until `\n` occurs)
112        //       See https://github.com/cryspen/hax/issues/1537.
113        for slice in stdout.split(b'\n') {
114            let msg = (|| {
115                let slice = slice.ok()?;
116                let mut de = serde_json::Deserializer::from_slice(&slice);
117                de.disable_recursion_limit();
118                let de = serde_stacker::Deserializer::new(&mut de);
119                let msg = ExtendedFromEngine::deserialize(de);
120                msg.ok()
121            })()
122            .expect(
123                "Hax engine sent an invalid json value. \
124                                This might be caused by debug messages on stdout, \
125                                which is reserved for JSON communication with cargo-hax",
126            );
127
128            match msg {
129                ExtendedFromEngine::Response(res) => response = Some(res),
130                ExtendedFromEngine::FromEngine(FromEngine::Exit) => break,
131                // Proxy messages from the OCaml engine
132                ExtendedFromEngine::FromEngine(from_engine) => {
133                    crate::hax_io::write(&from_engine);
134                    if from_engine.requires_response() {
135                        let ExtendedToEngine::ToEngine(response) = crate::hax_io::read() else {
136                            panic!(
137                                "The frontend sent an incorrect message: expected `ExtendedToEngine::ToEngine` since we sent a `ExtendedFromEngine::FromEngine`."
138                            )
139                        };
140                        send!(stdin, &response);
141                    }
142                }
143            }
144        }
145        drop(stdin);
146
147        let exit_status = engine_subprocess.wait().unwrap();
148        if !exit_status.success() {
149            panic!("ocaml engine crashed");
150        }
151
152        response
153    }
154}