1use std::path::Path;
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct Format {
7 pub id: String,
8 pub describe: String,
9 pub rule: Rule,
10}
11
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub enum Rule {
14 File { extensions: Vec<String>, magic: Option<Vec<u8>> },
15 Directory { files: Vec<String>, extensions: Vec<String> },
16}
17
18impl Format {
19 pub fn file(id: &str, describe: &str, extensions: &[&str], magic: Option<&[u8]>) -> Format {
20 Format {
21 id: id.into(),
22 describe: describe.into(),
23 rule: Rule::File {
24 extensions: extensions.iter().map(|e| e.to_string()).collect(),
25 magic: magic.map(<[u8]>::to_vec),
26 },
27 }
28 }
29
30 pub fn directory(id: &str, describe: &str, files: &[&str], extensions: &[&str]) -> Format {
31 Format {
32 id: id.into(),
33 describe: describe.into(),
34 rule: Rule::Directory {
35 files: files.iter().map(|f| f.to_string()).collect(),
36 extensions: extensions.iter().map(|e| e.to_string()).collect(),
37 },
38 }
39 }
40
41 pub fn matches(&self, path: &Path) -> bool {
42 match &self.rule {
43 Rule::File { extensions, magic } => {
44 let named = path
45 .extension()
46 .and_then(|e| e.to_str())
47 .is_some_and(|e| extensions.iter().any(|x| x.eq_ignore_ascii_case(e)));
48 named || magic.as_deref().is_some_and(|m| starts_with(path, m))
49 }
50 Rule::Directory { files, extensions } => {
51 path.is_dir()
52 && files.iter().all(|f| path.join(f).is_file())
53 && std::fs::read_dir(path)
54 .map(|rd| {
55 rd.flatten().any(|e| {
56 e.path()
57 .extension()
58 .and_then(|x| x.to_str())
59 .is_some_and(|x| extensions.iter().any(|want| want == x))
60 })
61 })
62 .unwrap_or(false)
63 }
64 }
65 }
66}
67
68fn starts_with(path: &Path, magic: &[u8]) -> bool {
69 if magic.is_empty() || !std::fs::metadata(path).is_ok_and(|m| m.is_file()) {
70 return false;
71 }
72 let mut head = vec![0u8; magic.len()];
73 std::fs::File::open(path)
74 .and_then(|mut f| std::io::Read::read_exact(&mut f, &mut head))
75 .map(|_| head == magic)
76 .unwrap_or(false)
77}
78
79#[derive(Debug, Clone, PartialEq, Eq)]
80pub enum TokenizerSource {
81 Library(String),
82 HuggingFace,
83 Engine,
84}
85
86pub fn valid_id(s: &str) -> bool {
87 !s.is_empty()
88 && !s.starts_with('.')
89 && s.chars().all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || matches!(c, '.' | '-' | '_'))
90}
91
92#[derive(Debug, Clone, PartialEq, Eq)]
93pub struct RuntimeInfo {
94 pub id: &'static str,
95 pub aliases: &'static [&'static str],
96 pub formats: Vec<Format>,
97 pub tokenizer: TokenizerSource,
98}
99
100#[derive(Debug, Clone, PartialEq, Eq)]
101pub struct Device {
102 pub backend: String,
103 pub name: String,
104 pub memory: u64,
105}
106
107impl Device {
108 pub fn summary(&self) -> String {
109 let mut s = format!("{}: {}", self.backend, self.name);
110 if self.memory > 0 {
111 s.push_str(&format!(", {:.1} GB", self.memory as f64 / 1e9));
112 }
113 s
114 }
115}
116
117#[cfg(test)]
118mod tests {
119 use super::*;
120 use std::path::PathBuf;
121
122 fn dir(tag: &str) -> PathBuf {
123 let d = std::env::temp_dir().join(format!("superfluid-artifact-{tag}-{}", std::process::id()));
124 let _ = std::fs::remove_dir_all(&d);
125 std::fs::create_dir_all(&d).unwrap();
126 d
127 }
128
129 #[test]
130 fn a_file_format_matches_by_extension_in_any_case_or_by_magic() {
131 let d = dir("file");
132 let gguf = Format::file("gguf", "a GGUF file", &["gguf"], Some(b"GGUF"));
133 std::fs::write(d.join("a.gguf"), b"").unwrap();
134 std::fs::write(d.join("a.GGUF"), b"").unwrap();
135 std::fs::write(d.join("weights.bin"), b"GGUF\x03\x00\x00\x00").unwrap();
136 std::fs::write(d.join("a.base"), b"BASE").unwrap();
137 assert!(gguf.matches(&d.join("a.gguf")));
138 assert!(gguf.matches(&d.join("a.GGUF")));
139 assert!(gguf.matches(&d.join("weights.bin")), "the magic names a file its name does not");
140 assert!(!gguf.matches(&d.join("a.base")));
141 assert!(!gguf.matches(&d.join("missing.bin")));
142 assert!(!gguf.matches(&d), "a directory has no magic");
143 let _ = std::fs::remove_dir_all(&d);
144 }
145
146 #[test]
147 fn a_directory_format_needs_its_files_and_one_weight() {
148 let d = dir("dir");
149 let mlx = Format::directory("mlx", "an MLX model directory", &["config.json"], &["safetensors"]);
150 assert!(!mlx.matches(&d));
151 std::fs::write(d.join("config.json"), b"{}").unwrap();
152 assert!(!mlx.matches(&d), "no weights");
153 std::fs::write(d.join("model.safetensors"), b"").unwrap();
154 assert!(mlx.matches(&d));
155 assert!(!mlx.matches(&d.join("config.json")), "a file is not a directory");
156 let _ = std::fs::remove_dir_all(&d);
157 }
158
159 #[test]
160 fn the_magic_probe_opens_only_regular_files() {
161 let d = dir("fifo");
162 let fifo = d.join("pipe");
163 let c = std::ffi::CString::new(fifo.to_string_lossy().as_bytes()).unwrap();
164 extern "C" {
165 fn mkfifo(path: *const std::ffi::c_char, mode: u32) -> i32;
166 }
167 assert_eq!(unsafe { mkfifo(c.as_ptr(), 0o600) }, 0, "mkfifo");
170 let gguf = Format::file("gguf", "a GGUF file", &["gguf"], Some(b"GGUF"));
171 assert!(!gguf.matches(&fifo));
172 let _ = std::fs::remove_dir_all(&d);
173 }
174
175 #[test]
176 fn a_device_sums_itself_up() {
177 let d = Device { backend: "Metal".into(), name: "Apple M5 Pro".into(), memory: 36_000_000_000 };
178 assert_eq!(d.summary(), "Metal: Apple M5 Pro, 36.0 GB");
179 let cpu = Device { backend: "CPU".into(), name: "Apple M5 Pro".into(), memory: 0 };
180 assert_eq!(cpu.summary(), "CPU: Apple M5 Pro");
181 }
182}