solar_cli/
utils.rs

1//! Utility functions used by the Solar CLI.
2
3use solar_interface::diagnostics::DiagCtxt;
4use std::io;
5
6#[cfg(feature = "tracing")]
7use solar_sema::ast::Either;
8
9#[cfg(feature = "mimalloc")]
10use mimalloc as _;
11#[cfg(all(feature = "jemalloc", unix))]
12use tikv_jemallocator as _;
13
14// Keep the system allocator in tests, where we spawn a ton of processes and any extra startup cost
15// slows down tests massively.
16cfg_if::cfg_if! {
17    if #[cfg(debug_assertions)] {
18        type AllocatorInner = std::alloc::System;
19    } else if #[cfg(feature = "mimalloc")] {
20        type AllocatorInner = mimalloc::MiMalloc;
21    } else if #[cfg(all(feature = "jemalloc", unix))] {
22        type AllocatorInner = tikv_jemallocator::Jemalloc;
23    } else {
24        type AllocatorInner = std::alloc::System;
25    }
26}
27
28cfg_if::cfg_if! {
29    if #[cfg(feature = "tracy-allocator")] {
30        pub(super) type WrappedAllocator = tracing_tracy::client::ProfiledAllocator<AllocatorInner>;
31        pub(super) const fn new_wrapped_allocator() -> WrappedAllocator {
32            Allocator::new(AllocatorInner {}, 100)
33        }
34    } else {
35        pub(super) type WrappedAllocator = AllocatorInner;
36        pub(super) const fn new_wrapped_allocator() -> WrappedAllocator {
37            AllocatorInner {}
38        }
39    }
40}
41
42/// The global allocator used by the compiler.
43pub type Allocator = WrappedAllocator;
44
45/// Create a new instance of the global allocator.
46pub const fn new_allocator() -> Allocator {
47    new_wrapped_allocator()
48}
49
50/// `tracing` logger destination.
51#[derive(Default)]
52pub enum LogDestination {
53    /// [`io::stdout`].
54    #[default]
55    Stdout,
56    /// [`io::stderr`].
57    Stderr,
58}
59
60#[cfg(feature = "tracing")]
61impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for LogDestination {
62    type Writer = Either<io::Stdout, io::Stderr>;
63
64    fn make_writer(&'a self) -> Self::Writer {
65        match self {
66            Self::Stdout => Either::Left(io::stdout()),
67            Self::Stderr => Either::Right(io::stderr()),
68        }
69    }
70}
71
72/// Initialize the tracing logger.
73#[must_use]
74pub fn init_logger(dst: LogDestination) -> impl Sized {
75    #[cfg(not(feature = "tracing"))]
76    {
77        if std::env::var_os("RUST_LOG").is_some() {
78            let msg = "`RUST_LOG` is set, but \"tracing\" support was not enabled at compile time";
79            DiagCtxt::new_early().warn(msg).emit();
80        }
81        if std::env::var_os("SOLAR_PROFILE").is_some() {
82            let msg =
83                "`SOLAR_PROFILE` is set, but \"tracing\" support was not enabled at compile time";
84            DiagCtxt::new_early().warn(msg).emit();
85        }
86    }
87
88    #[cfg(feature = "tracing")]
89    match try_init_logger(dst) {
90        Ok(guard) => guard,
91        Err(e) => DiagCtxt::new_early().fatal(e).emit(),
92    }
93}
94
95#[cfg(feature = "tracing")]
96fn try_init_logger(dst: LogDestination) -> Result<impl Sized, String> {
97    use tracing_subscriber::prelude::*;
98
99    let (profile_layer, guard) = match std::env::var("SOLAR_PROFILE").as_deref() {
100        Ok("chrome") => {
101            if !cfg!(feature = "tracing-chrome") {
102                return Err("chrome profiler support is not compiled in".to_string());
103            }
104            let (layer, guard) = chrome_layer();
105            (Some(layer.boxed()), Some(guard))
106        }
107        Ok("tracy") => {
108            if !cfg!(feature = "tracy") {
109                return Err("tracy profiler support is not compiled in".to_string());
110            }
111            (Some(tracy_layer().boxed()), Default::default())
112        }
113        Ok(s) => return Err(format!("unknown profiler '{s}'; valid values: 'chrome', 'tracy'")),
114        Err(_) => Default::default(),
115    };
116    tracing_subscriber::Registry::default()
117        .with(tracing_subscriber::EnvFilter::from_default_env())
118        .with(profile_layer)
119        .with(tracing_subscriber::fmt::layer().with_writer(dst))
120        .try_init()
121        .map(|()| guard)
122        .map_err(|e| e.to_string())
123}
124
125#[cfg(feature = "tracing")]
126#[cfg(feature = "tracy")]
127fn tracy_layer() -> tracing_tracy::TracyLayer<impl tracing_tracy::Config> {
128    struct Config(tracing_subscriber::fmt::format::DefaultFields);
129    impl tracing_tracy::Config for Config {
130        type Formatter = tracing_subscriber::fmt::format::DefaultFields;
131        fn formatter(&self) -> &Self::Formatter {
132            &self.0
133        }
134        fn format_fields_in_zone_name(&self) -> bool {
135            false
136        }
137    }
138
139    tracing_tracy::client::register_demangler!();
140
141    tracing_tracy::TracyLayer::new(Config(Default::default()))
142}
143
144#[cfg(feature = "tracing")]
145#[cfg(not(feature = "tracy"))]
146fn tracy_layer() -> tracing_subscriber::layer::Identity {
147    tracing_subscriber::layer::Identity::new()
148}
149
150#[cfg(feature = "tracing")]
151#[cfg(feature = "tracing-chrome")]
152fn chrome_layer<S>() -> (tracing_chrome::ChromeLayer<S>, tracing_chrome::FlushGuard)
153where
154    S: tracing::Subscriber
155        + for<'span> tracing_subscriber::registry::LookupSpan<'span>
156        + Send
157        + Sync,
158{
159    tracing_chrome::ChromeLayerBuilder::new().include_args(true).build()
160}
161
162#[cfg(feature = "tracing")]
163#[cfg(not(feature = "tracing-chrome"))]
164fn chrome_layer() -> (tracing_subscriber::layer::Identity, ()) {
165    (tracing_subscriber::layer::Identity::new(), ())
166}
167
168/*
169pub(crate) fn env_to_bool(value: Option<&std::ffi::OsStr>) -> bool {
170    value.is_some_and(|value| value == "1" || value == "true")
171}
172*/