Skip to main content

teeny_compiler/compiler/
mod.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use std::path::PathBuf;
18use std::process::Command;
19
20use anyhow::Context;
21
22use crate::errors::Result;
23
24/// Compilation backends (LLVM/MLIR, `ndarray`).
25pub mod backend;
26/// Compiler drivers (invoking `teenyc` and friends per target).
27pub mod driver;
28/// Compilation target descriptions (device capability, etc).
29pub mod target;
30
31/// Resolve the `teenyc` binary to invoke.
32///
33/// Priority:
34/// 1. `$TEENYC_PATH`, if set.
35/// 2. The sole `rustup`-linked toolchain whose name contains `teenyc` — the naming convention
36///    `cargo teeny install-toolchain` uses (default toolchain name is `<channel>-<host>` with
37///    `channel` defaulting to `stable-teenyc`; see `cargo-teeny`'s `install_toolchain` module) —
38///    resolved to a binary path via `rustup which --toolchain <name> teenyc`.
39///
40/// This deliberately does not fall back further to a bare `teenyc` looked up on `$PATH`: that
41/// would only work by accident (most `$PATH`s don't have a `teenyc` on them at all) and produces
42/// a worse error than pointing at the two supported setup paths.
43pub fn find_teenyc() -> Result<PathBuf> {
44    if let Ok(path) = std::env::var("TEENYC_PATH") {
45        return Ok(PathBuf::from(path));
46    }
47
48    let toolchain = sole_teenyc_toolchain()?;
49    which_in_toolchain(&toolchain)
50}
51
52/// Names of installed `rustup` toolchains containing `teenyc`, parsed from `rustup toolchain
53/// list` output (one name per line, optionally suffixed with ` (default)`/` (active)`).
54fn teenyc_toolchain_names(rustup_toolchain_list_output: &str) -> Vec<String> {
55    rustup_toolchain_list_output
56        .lines()
57        .filter_map(|line| line.split_whitespace().next())
58        .filter(|name| name.contains("teenyc"))
59        .map(str::to_string)
60        .collect()
61}
62
63/// The single installed `rustup` toolchain whose name contains `teenyc`. Errors if none or more
64/// than one is found — in the latter case the caller needs `TEENYC_PATH` to disambiguate.
65fn sole_teenyc_toolchain() -> Result<String> {
66    let output = Command::new("rustup")
67        .args(["toolchain", "list"])
68        .output()
69        .context("spawn `rustup toolchain list` (is rustup installed and on PATH?)")?;
70    anyhow::ensure!(
71        output.status.success(),
72        "`rustup toolchain list` exited with {}",
73        output.status
74    );
75
76    let names = teenyc_toolchain_names(&String::from_utf8_lossy(&output.stdout));
77    match names.as_slice() {
78        [] => anyhow::bail!(
79            "no teenyc rustup toolchain found; set TEENYC_PATH to the teenyc binary, or install \
80             one with `cargo teeny install-toolchain` (see cargo-teeny)"
81        ),
82        [name] => Ok(name.clone()),
83        multiple => anyhow::bail!(
84            "multiple teenyc rustup toolchains found ({}); set TEENYC_PATH to disambiguate",
85            multiple.join(", ")
86        ),
87    }
88}
89
90/// Resolves `toolchain`'s `teenyc` binary path via `rustup which --toolchain <toolchain> teenyc`.
91fn which_in_toolchain(toolchain: &str) -> Result<PathBuf> {
92    let output = Command::new("rustup")
93        .args(["which", "--toolchain", toolchain, "teenyc"])
94        .output()
95        .with_context(|| format!("spawn `rustup which --toolchain {toolchain} teenyc`"))?;
96    anyhow::ensure!(
97        output.status.success(),
98        "`rustup which --toolchain {toolchain} teenyc` exited with {}",
99        output.status
100    );
101    Ok(PathBuf::from(
102        String::from_utf8_lossy(&output.stdout).trim().to_string(),
103    ))
104}
105
106/// Resolve the effective kernel cache directory.
107///
108/// Priority: `$TEENYC_CACHE_DIR` (if set) > `<exe_dir>/../cache` (if that
109/// directory exists — the layout `cargo teeny package` produces, with
110/// `cache/` sitting next to `bin/`) > `/tmp/teenyc_cache`.
111///
112/// The exe-relative check only ever fires when a real `cache/` directory is
113/// actually there, so plain `cargo run`/`cargo test` dev builds (whose exe
114/// lives under `target/debug/...`, with no `cache/` sibling) are unaffected.
115pub fn default_cache_dir() -> String {
116    if let Ok(dir) = std::env::var("TEENYC_CACHE_DIR") {
117        return dir;
118    }
119
120    match std::env::current_exe() {
121        Ok(exe) => sibling_cache_dir(&exe).unwrap_or_else(|| "/tmp/teenyc_cache".to_string()),
122        Err(_) => "/tmp/teenyc_cache".to_string(),
123    }
124}
125
126/// `<exe's parent's parent>/cache`, if that directory exists.
127fn sibling_cache_dir(exe: &std::path::Path) -> Option<String> {
128    let package_root = exe.parent()?.parent()?;
129    let candidate = package_root.join("cache");
130    candidate
131        .is_dir()
132        .then(|| candidate.to_string_lossy().into_owned())
133}
134
135#[cfg(test)]
136mod find_teenyc_tests {
137    use super::*;
138
139    #[test]
140    fn finds_single_teenyc_toolchain() {
141        let output = "stable-x86_64-unknown-linux-gnu (default)\n\
142                       stable-teenyc-x86_64-unknown-linux-gnu\n";
143        assert_eq!(
144            teenyc_toolchain_names(output),
145            vec!["stable-teenyc-x86_64-unknown-linux-gnu"]
146        );
147    }
148
149    #[test]
150    fn empty_when_no_teenyc_toolchain() {
151        let output =
152            "stable-x86_64-unknown-linux-gnu (default)\nnightly-x86_64-unknown-linux-gnu\n";
153        assert!(teenyc_toolchain_names(output).is_empty());
154    }
155
156    #[test]
157    fn finds_multiple_teenyc_toolchains() {
158        let output = "stable-teenyc-x86_64-unknown-linux-gnu (default)\n\
159                       my-teenyc-toolchain\n";
160        assert_eq!(
161            teenyc_toolchain_names(output),
162            vec![
163                "stable-teenyc-x86_64-unknown-linux-gnu",
164                "my-teenyc-toolchain"
165            ]
166        );
167    }
168}
169
170#[cfg(test)]
171mod cache_dir_tests {
172    use std::fs;
173    use std::time::{SystemTime, UNIX_EPOCH};
174
175    use super::*;
176
177    fn tmp_root(name: &str) -> std::path::PathBuf {
178        let suffix = SystemTime::now()
179            .duration_since(UNIX_EPOCH)
180            .unwrap()
181            .as_nanos();
182        std::env::temp_dir().join(format!("teeny-compiler-cache-dir-test-{name}-{suffix}"))
183    }
184
185    #[test]
186    fn finds_sibling_cache_dir_when_present() {
187        let root = tmp_root("present");
188        let bin_dir = root.join("bin");
189        fs::create_dir_all(&bin_dir).unwrap();
190        fs::create_dir_all(root.join("cache")).unwrap();
191        let exe = bin_dir.join("myapp");
192
193        let found = sibling_cache_dir(&exe).expect("cache dir should be found");
194        assert_eq!(found, root.join("cache").to_string_lossy());
195
196        let _ = fs::remove_dir_all(&root);
197    }
198
199    #[test]
200    fn none_when_cache_dir_missing() {
201        let root = tmp_root("missing");
202        let bin_dir = root.join("bin");
203        fs::create_dir_all(&bin_dir).unwrap();
204        let exe = bin_dir.join("myapp");
205
206        assert!(sibling_cache_dir(&exe).is_none());
207
208        let _ = fs::remove_dir_all(&root);
209    }
210
211    #[test]
212    fn none_when_exe_has_no_grandparent() {
213        assert!(sibling_cache_dir(std::path::Path::new("myapp")).is_none());
214    }
215}