Skip to main content

micro_wakeword/
runtime.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use tflite_c::TfLiteLibrary;
5
6use crate::{Error, Result};
7
8/// Selects the TensorFlow Lite C runtime used for inference.
9#[derive(Clone, Debug, Default)]
10pub enum Runtime {
11    /// Check explicit environment overrides, then use the bundled runtime.
12    #[default]
13    Auto,
14    /// Load a specific shared library.
15    Path(PathBuf),
16    /// Search standard system locations supported by `tflite-c-rs`.
17    System,
18}
19
20impl Runtime {
21    pub fn from_path(path: impl Into<PathBuf>) -> Self {
22        Self::Path(path.into())
23    }
24
25    pub(crate) fn load(&self) -> Result<Arc<TfLiteLibrary>> {
26        let explicit = match self {
27            Self::Path(path) => Some(path.clone()),
28            Self::Auto => std::env::var_os("MICRO_WAKEWORD_TFLITE_LIB")
29                .or_else(|| std::env::var_os("TFLITE_C_LIB"))
30                .map(PathBuf::from),
31            Self::System => None,
32        };
33        if let Some(path) = explicit {
34            return TfLiteLibrary::load_from_path(&path).map_err(Error::from);
35        }
36        match self {
37            Self::Auto => {
38                let path = bundled_runtime_path()?;
39                TfLiteLibrary::load_from_path(path).map_err(Error::from)
40            }
41            Self::System => TfLiteLibrary::load_default().map_err(Error::from),
42            Self::Path(_) => unreachable!(),
43        }
44    }
45
46    pub fn path(&self) -> Option<&Path> {
47        match self {
48            Self::Path(path) => Some(path),
49            Self::Auto | Self::System => None,
50        }
51    }
52}
53
54#[cfg(all(target_os = "windows", target_arch = "x86_64"))]
55fn bundled_runtime_path() -> Result<PathBuf> {
56    use std::fs;
57    use std::io::Write;
58
59    const BYTES: &[u8] = include_bytes!("../runtime/windows-x86_64/tensorflowlite_c-2.17.1.dll");
60    const SHA256: &str = "882e6d8f9866ff84f23d4b964c145b7f0f0a8907fa830dcd8c499e7c46bf3365";
61    let local = std::env::var_os("LOCALAPPDATA").ok_or_else(|| {
62        Error::UnsupportedPlatform("LOCALAPPDATA is unavailable for runtime extraction".into())
63    })?;
64    let directory = PathBuf::from(local)
65        .join("micro-wakeword")
66        .join("runtime-2.17.1");
67    let destination = directory.join("tensorflowlite_c.dll");
68    if destination.exists() && file_sha256(&destination)? == SHA256 {
69        return Ok(destination);
70    }
71
72    fs::create_dir_all(&directory).map_err(|source| Error::Io {
73        path: directory.clone(),
74        source,
75    })?;
76    let temporary = directory.join(format!("tensorflowlite_c.{}.tmp", std::process::id()));
77    let mut file = fs::File::create(&temporary).map_err(|source| Error::Io {
78        path: temporary.clone(),
79        source,
80    })?;
81    file.write_all(BYTES)
82        .and_then(|_| file.sync_all())
83        .map_err(|source| Error::Io {
84            path: temporary.clone(),
85            source,
86        })?;
87    if file_sha256(&temporary)? != SHA256 {
88        return Err(Error::InvalidConfig(
89            "embedded TensorFlow Lite runtime checksum mismatch".into(),
90        ));
91    }
92    if destination.exists() {
93        fs::remove_file(&destination).map_err(|source| Error::Io {
94            path: destination.clone(),
95            source,
96        })?;
97    }
98    if let Err(source) = fs::rename(&temporary, &destination) {
99        if destination.exists() && file_sha256(&destination)? == SHA256 {
100            let _ = fs::remove_file(&temporary);
101        } else {
102            return Err(Error::Io {
103                path: destination,
104                source,
105            });
106        }
107    }
108    Ok(destination)
109}
110
111#[cfg(all(target_os = "windows", target_arch = "x86_64"))]
112fn file_sha256(path: &Path) -> Result<String> {
113    use sha2::{Digest, Sha256};
114    let bytes = std::fs::read(path).map_err(|source| Error::Io {
115        path: path.to_owned(),
116        source,
117    })?;
118    Ok(format!("{:x}", Sha256::digest(bytes)))
119}
120
121#[cfg(not(all(target_os = "windows", target_arch = "x86_64")))]
122fn bundled_runtime_path() -> Result<PathBuf> {
123    Err(Error::UnsupportedPlatform(
124        "the bundled runtime currently supports only Windows x86-64; use Runtime::Path or Runtime::System".into(),
125    ))
126}