implement BridgeStan download and module compilation on Rust

.github/workflows/main.yaml

@@ -3,7 +3,7 @@ name: bridgestan tests
 on:
   push:
     branches:
-      - 'main'
+      - "main"
   pull_request:
   workflow_dispatch: {}
 
@@ -128,7 +128,6 @@ jobs:
         env:
           BRIDGESTAN: ${{ github.workspace }}
 
-
   julia:
     needs: [build]
     runs-on: ${{matrix.os}}
@@ -290,5 +289,5 @@ jobs:
         run: |
           cargo clippy
           cargo fmt --check
-          cargo run --example=example
-          cargo test --verbose
+          cargo run --example=example --features compile-stan-model
+          cargo test --verbose --features compile-stan-model -- --test-threads=1

docs/languages/rust.rst

@@ -14,13 +14,9 @@ The BridgeStan Rust client is available on `crates.io <https://crates.io/crates/
 
     cargo add bridgestan
 
-To build and use BridgeStan models, a copy of the BridgeStan C++ source code
-is required. Please follow the :doc:`Getting Started guide <../getting-started>`
-or use the Rust client in tandem with an interface such as :doc:`Python <./python>`
-which automates this process.
+The first time you compile a model, the BridgeStan source code will be downloaded to `~/.bridgestan`. If you prefer to use a source distribution of BridgeStan, you can pass its path as the `bs_path` argument to `compile_model`.
 
-``STAN_THREADS=true`` needs to be specified when compiling a model, for more
-details see the `API reference <https://docs.rs/bridgestan>`__.
+Note that the system pre-requisites from the [Getting Started Guide](../getting-started.rst) are still required and will not be automatically installed by this method.
 
 Example Program
 ---------------

rust/.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+    "rust-analyzer.cargo.features": [
+        "compile-stan-model"
+    ]
+}

rust/Cargo.toml

@@ -11,10 +11,22 @@ homepage = "https://roualdes.github.io/bridgestan/latest/"
 [dependencies]
 libloading = "0.8.0"
 thiserror = "1.0.40"
+ureq = { version = "2.7", optional = true }
+tar = { version = "0.4", optional = true }
+flate2 = { version = "1.0", optional = true }
+dirs = { version = "5.0", optional = true }
+path-absolutize = { version = "3.1", optional = true }
+
+[features]
+compile-stan-model = ["ureq", "tar", "flate2", "dirs", "path-absolutize"]
 
 [build-dependencies]
 bindgen = "0.69.1"
 
 [dev-dependencies]
 approx = "0.5.1"
 rand = "0.8.5"
+
+[[example]]
+name = "example"
+required-features = ["compile-stan-model"]

rust/README.md

@@ -9,44 +9,24 @@ natively from Rust.
 Internally, it relies on [`bindgen`](https://docs.rs/bindgen/) and
 [`libloading`](https://docs.rs/libloading/).
 
-## Compiling the model
-
-The Rust wrapper does not currently have any functionality to compile Stan models.
-Compiled shared libraries need to be built manually using `make` or with the Julia
-or Python bindings.
-
-For safety reasons all Stan models need to be installed with `STAN_THREADS=true`.
-When compiling a model using `make`, set the environment variable:
-
-```bash
-STAN_THREADS=true make some_model
-```
-
-When compiling a Stan model in python, this has to be specified in the `make_args`
-argument:
-
-```python
-path = bridgestan.compile_model("stan_model.stan", make_args=["STAN_THREADS=true"])
-```
-
-If `STAN_THREADS` was not specified while building the model, the Rust wrapper
-will throw an error when loading the model.
-
-## Usage:
+## Usage
 
 Run this example with `cargo run --example=example`.
 
 ```rust
 use std::ffi::CString;
 use std::path::Path;
-use bridgestan::{BridgeStanError, Model, open_library};
+use bridgestan::{BridgeStanError, Model, open_library, compile_model};
 
-// The path to the compiled model.
-// Get for instance from python `bridgestan.compile_model`
+// The path to the Stan model
 let path = Path::new(env!["CARGO_MANIFEST_DIR"])
     .parent()
     .unwrap()
-    .join("test_models/simple/simple_model.so");
+    .join("test_models/simple/simple.stan");
+
+// The path to the compiled model
+let path = compile_model(path, vec![], vec![], None).expect("Could not compile Stan model.");
+println!("Compiled model: {:?}", path);
 
 let lib = open_library(path).expect("Could not load compiled Stan model.");
 

rust/examples/example.rs

@@ -1,14 +1,19 @@
-use bridgestan::{open_library, BridgeStanError, Model};
+use bridgestan::{compile_model, open_library, BridgeStanError, Model};
 use std::ffi::CString;
 use std::path::Path;
 
 fn main() {
-    // The path to the compiled model.
-    // Get for instance from python `bridgestan.compile_model`
+    // The path to the Stan model
     let path = Path::new(env!["CARGO_MANIFEST_DIR"])
         .parent()
         .unwrap()
-        .join("test_models/simple/simple_model.so");
+        .join("test_models")
+        .join("simple")
+        .join("simple.stan");
+
+    // The path to the compiled model
+    let path = compile_model(path, vec![], vec![], None).expect("Could not compile Stan model.");
+    println!("Compiled model: {:?}", path);
 
     let lib = open_library(path).expect("Could not load compiled Stan model.");
 

rust/src/bs_safe.rs

@@ -1,3 +1,5 @@
+#[cfg(feature = "compile-stan-model")]
+use crate::download_compile::VERSION;
 use crate::ffi;
 use std::borrow::Borrow;
 use std::collections::hash_map::DefaultHasher;
@@ -101,9 +103,17 @@ pub enum BridgeStanError {
     /// Setting a print-callback failed.
     #[error("Failed to set a print-callback: {0}")]
     SetCallbackFailed(String),
+    #[cfg(feature = "compile-stan-model")]
+    /// Setting a compile Stan model failed.
+    #[error("Failed to compile Stan model: {0}")]
+    ModelCompilingFailed(String),
+    #[cfg(feature = "compile-stan-model")]
+    /// Setting a download BridgeStan failed.
+    #[error("Failed to download BridgeStan {VERSION} from github.com: {0}")]
+    DownloadFailed(String),
 }
 
-type Result<T> = std::result::Result<T, BridgeStanError>;
+pub(crate) type Result<T> = std::result::Result<T, BridgeStanError>;
 
 /// Open a compiled Stan library.
 ///

rust/src/download_compile.rs

@@ -0,0 +1,133 @@
+use crate::bs_safe::{BridgeStanError, Result};
+use flate2::read::GzDecoder;
+use path_absolutize::Absolutize;
+use std::{
+    env::temp_dir,
+    fs,
+    path::{Path, PathBuf},
+};
+use tar::Archive;
+
+pub(crate) const VERSION: &str = env!("CARGO_PKG_VERSION");
+
+/// Download and unzip the BridgeStan source distribution for this version
+/// to ~/.bridgestan/bridgestan-version
+pub fn get_bridgestan_src() -> Result<PathBuf> {
+    let homedir = dirs::home_dir().unwrap_or(temp_dir());
+
+    let bs_path_download_temp = homedir.join(".bridgestan_tmp_dir");
+    let bs_path_download = homedir.join(".bridgestan");
+
+    let bs_path_download_temp_join_version =
+        bs_path_download_temp.join(format!("bridgestan-{VERSION}"));
+    let bs_path_download_join_version = bs_path_download.join(format!("bridgestan-{VERSION}"));
+
+    if !bs_path_download_join_version.exists() {
+        println!("Downloading BridgeStan");
+
+        fs::remove_dir_all(&bs_path_download_temp).unwrap_or_default();
+        fs::create_dir(&bs_path_download_temp).unwrap_or_default();
+        fs::create_dir(&bs_path_download).unwrap_or_default();
+
+        let url = "https://github.com/roualdes/bridgestan/releases/download/".to_owned()
+            + format!("v{VERSION}/bridgestan-{VERSION}.tar.gz").as_str();
+
+        let response = ureq::get(url.as_str())
+            .call()
+            .map_err(|e| BridgeStanError::DownloadFailed(e.to_string()))?;
+        let len = response
+            .header("Content-Length")
+            .and_then(|s| s.parse::<usize>().ok())
+            .unwrap_or(50_000_000);
+
+        let mut bytes: Vec<u8> = Vec::with_capacity(len);
+        response
+            .into_reader()
+            .read_to_end(&mut bytes)
+            .map_err(|e| BridgeStanError::DownloadFailed(e.to_string()))?;
+
+        let tar = GzDecoder::new(bytes.as_slice());
+        let mut archive = Archive::new(tar);
+        archive
+            .unpack(&bs_path_download_temp)
+            .map_err(|e| BridgeStanError::DownloadFailed(e.to_string()))?;
+
+        fs::rename(
+            bs_path_download_temp_join_version,
+            &bs_path_download_join_version,
+        )
+        .map_err(|e| BridgeStanError::DownloadFailed(e.to_string()))?;
+
+        fs::remove_dir(bs_path_download_temp).unwrap_or_default();
+
+        println!("Finished downloading BridgeStan");
+    }
+
+    Ok(bs_path_download_join_version)
+}
+
+/// Compile a Stan Model given a stan_file and the path to BridgeStan
+/// if None, then calls get_bridgestan_src() to download BridgeStan
+pub fn compile_model<P>(
+    stan_file: P,
+    stanc_args: Vec<&str>,
+    make_args: Vec<&str>,
+    bs_path: Option<P>,
+) -> Result<PathBuf>
+where
+    P: AsRef<Path>,
+{
+    let bs_path = match bs_path {
+        Some(path) => path.as_ref().to_owned(),
+        None => get_bridgestan_src()?,
+    };
+
+    let stan_file = stan_file
+        .as_ref()
+        .absolutize()
+        .map_err(|e| BridgeStanError::ModelCompilingFailed(e.to_string()))?;
+
+    if stan_file.extension().unwrap_or_default() != "stan" {
+        return Err(BridgeStanError::ModelCompilingFailed(
+            "File must be a .stan file".to_owned(),
+        ));
+    }
+
+    // add _model suffix and change extension to .so
+    let output = stan_file.with_extension("");
+    let output = output.with_file_name(format!(
+        "{}_model",
+        output.file_name().unwrap_or_default().to_string_lossy()
+    ));
+    let output = output.with_extension("so");
+
+    let stanc_args = [["--include-paths=."].as_slice(), stanc_args.as_slice()].concat();
+    let stanc_args = stanc_args.join(" ");
+    let stanc_args = format!("STANCFLAGS={}", stanc_args);
+    let stanc_args = [stanc_args.as_str()];
+
+    let cmd = [
+        &[output.to_str().unwrap_or_default()],
+        make_args.as_slice(),
+        stanc_args.as_slice(),
+    ]
+    .concat();
+
+    println!("Compiling model");
+    let proc = std::process::Command::new("make")
+        .args(cmd)
+        .current_dir(bs_path)
+        .env("STAN_THREADS", "true")
+        .output()
+        .map_err(|e| BridgeStanError::ModelCompilingFailed(e.to_string()))?;
+    println!("Finished compiling model");
+
+    if !proc.status.success() {
+        return Err(BridgeStanError::ModelCompilingFailed(format!(
+            "{} {}",
+            String::from_utf8_lossy(proc.stdout.as_slice()).into_owned(),
+            String::from_utf8_lossy(proc.stderr.as_slice()).into_owned(),
+        )));
+    }
+    Ok(output)
+}

rust/src/lib.rs

@@ -1,6 +1,12 @@
 #![doc = include_str!("../README.md")]
 
 mod bs_safe;
-pub(crate) mod ffi;
 
+#[cfg(feature = "compile-stan-model")]
+mod download_compile;
+
+pub(crate) mod ffi;
 pub use bs_safe::{open_library, BridgeStanError, Model, Rng, StanLibrary};
+
+#[cfg(feature = "compile-stan-model")]
+pub use download_compile::{compile_model, get_bridgestan_src};

rust/tests/model.rs

@@ -1,11 +1,11 @@
 mod common;
-use std::{f64::consts::PI, ffi::CString};
+use std::{f64::consts::PI, ffi::CString, fs::remove_file};
 
-use common::get_model;
+use common::{get_model, model_dir};
 
 use approx::{assert_abs_diff_eq, assert_ulps_eq};
 
-use bridgestan::{BridgeStanError, Model};
+use bridgestan::{compile_model, BridgeStanError, Model};
 
 #[test]
 fn throw_data() {
@@ -25,17 +25,39 @@ fn throw_data() {
 fn bad_arglength() {
     let (lib, data) = get_model("stdnormal");
     let model = Model::new(&lib, data, 42).unwrap();
-    let theta = vec![];
-    let mut grad = vec![];
+    let theta = [];
+    let mut grad = [];
     let _ = model.log_density_gradient(&theta[..], true, true, &mut grad[..]);
 }
 
 #[test]
 fn logp_gradient() {
     let (lib, data) = get_model("stdnormal");
     let model = Model::new(&lib, data, 42).unwrap();
-    let theta = vec![1f64];
-    let mut grad = vec![0f64];
+    let theta = [1f64];
+    let mut grad = [0f64];
+    let logp = model
+        .log_density_gradient(&theta[..], false, true, &mut grad[..])
+        .unwrap();
+    assert_ulps_eq!(logp, (2. * PI).sqrt().recip().ln() - 0.5);
+    assert_ulps_eq!(grad[0], -1f64);
+}
+
+#[cfg(target_family = "unix")]
+#[test]
+fn model_compiling() {
+    let name = "stdnormal";
+    let mut base = model_dir();
+    base.push(name);
+    let lib_path = base.join(format!("{}_model.so", name));
+    let stan_path = base.join(format!("{}.stan", name));
+    remove_file(lib_path).unwrap();
+    compile_model(stan_path, vec![], vec![], None).unwrap();
+
+    let (lib, data) = get_model(name);
+    let model = Model::new(&lib, data, 42).unwrap();
+    let theta = [1f64];
+    let mut grad = [0f64];
     let logp = model
         .log_density_gradient(&theta[..], false, true, &mut grad[..])
         .unwrap();
@@ -47,8 +69,8 @@ fn logp_gradient() {
 fn logp_hessian() {
     let (lib, data) = get_model("stdnormal");
     let model = Model::new(&lib, data, 42).unwrap();
-    let theta = vec![1f64];
-    let mut grad = vec![0f64];
+    let theta = [1f64];
+    let mut grad = [0f64];
     let mut hessian = vec![0f64];
     let logp = model
         .log_density_hessian(&theta[..], false, true, &mut grad[..], &mut hessian)