ext/
utils.rs

1//! A module containing various utility functions related to user interaction in some way.
2use std::{path::PathBuf, sync::Arc};
3
4use algebra::{
5    AlgebraType, MilnorAlgebra, SteenrodAlgebra,
6    module::{FDModule, Module, SteenrodModule, steenrod_module},
7};
8use anyhow::{Context, anyhow};
9use serde_json::Value;
10use sseq::coordinates::{Bidegree, BidegreeGenerator};
11
12use crate::{
13    CCC,
14    chain_complex::{AugmentedChainComplex, BoundedChainComplex, ChainComplex, FiniteChainComplex},
15    resolution::{Resolution, UnstableResolution},
16    save::SaveDirectory,
17};
18
19// We build docs with --all-features so the docs are at the feature = "nassau" version
20#[cfg(not(feature = "nassau"))]
21pub type QueryModuleResolution = Resolution<CCC>;
22
23/// The type returned by [`query_module`]. The value of this type depends on whether
24/// [`nassau`](crate::nassau) is enabled. In any case, it is an augmented free chain complex over
25/// either [`SteenrodAlgebra`] or [`MilnorAlgebra`] and supports the `compute_through_stem`
26/// function.
27#[cfg(feature = "nassau")]
28pub type QueryModuleResolution = crate::nassau::Resolution<FDModule<MilnorAlgebra>>;
29
30const STATIC_MODULES_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../ext/steenrod_modules");
31
32/// A config object is an object that specifies how a Steenrod module should be constructed.
33#[derive(Clone, Debug, Eq, PartialEq)]
34pub struct Config {
35    /// The json specification of the module
36    module: Value,
37    /// The basis for the Steenrod algebra
38    algebra: AlgebraType,
39}
40
41/// Given a module specification string, load a json description of the module as described
42/// [here](../index.html#module-specification).
43pub fn parse_module_name(module_name: &str) -> anyhow::Result<Value> {
44    let mut args = module_name.split('[');
45    let module_file = args.next().unwrap();
46    let mut module = load_module_json(module_file)
47        .with_context(|| format!("Failed to load module file {module_file}"))?;
48    if let Some(shift) = args.next() {
49        let shift: i64 = match shift.strip_suffix(']') {
50            None => return Err(anyhow!("Unterminated shift [")),
51            Some(x) => x
52                .parse()
53                .with_context(|| format!("Cannot parse shift value ({x}) as an integer"))?,
54        };
55        if let Some(spec_shift) = module.get_mut("shift") {
56            *spec_shift = Value::from(spec_shift.as_i64().unwrap() + shift);
57        } else {
58            module["shift"] = Value::from(shift);
59        }
60    }
61    Ok(module)
62}
63
64impl TryFrom<&str> for Config {
65    type Error = anyhow::Error;
66
67    fn try_from(spec: &str) -> Result<Self, Self::Error> {
68        let mut args = spec.split('@');
69        let module_name = args.next().unwrap();
70        let algebra = match args.next() {
71            Some(x) => x
72                .parse()
73                .with_context(|| format!("Invalid algebra type: {x}"))?,
74            None => AlgebraType::Milnor,
75        };
76
77        Ok(Self {
78            module: parse_module_name(module_name)
79                .with_context(|| format!("Failed to load module: {module_name}"))?,
80            algebra,
81        })
82    }
83}
84
85impl<T, E> TryFrom<(&str, T)> for Config
86where
87    anyhow::Error: From<E>,
88    T: TryInto<AlgebraType, Error = E>,
89{
90    type Error = anyhow::Error;
91
92    fn try_from(mut spec: (&str, T)) -> Result<Self, Self::Error> {
93        let algebra = spec.1.try_into()?;
94        if spec.0.contains('@') {
95            if spec.0.ends_with(&*algebra.to_string()) {
96                spec.0 = &spec.0[0..spec.0.len() - algebra.to_string().len() - 1];
97            } else {
98                return Err(anyhow!("Invalid algebra supplied. Must be {}", algebra));
99            }
100        }
101        Ok(Self {
102            module: parse_module_name(spec.0)?,
103            algebra,
104        })
105    }
106}
107
108impl<T: TryInto<AlgebraType>> TryFrom<(Value, T)> for Config {
109    type Error = T::Error;
110
111    fn try_from(spec: (Value, T)) -> Result<Self, Self::Error> {
112        Ok(Self {
113            module: spec.0,
114            algebra: spec.1.try_into()?,
115        })
116    }
117}
118
119/// This constructs a resolution resolving a module according to the specifications
120///
121/// # Arguments
122///  - `module_spec`: A specification for the module. This is any object that implements
123///    [`TryInto<Config>`] (with appropriate error bounds). In practice, we can supply
124///    - A [`Config`] object itself
125///    - `(json, algebra)`: The first argument is a [`serde_json::Value`] that specifies the
126///      module; the second argument is either a string (`"milnor"` or `"adem"`) or an
127///      [`algebra::AlgebraType`] object.
128///    - `(module_name, algebra)`: The first argument is the name of the module and the second is
129///      as above. Modules are searched in the current directory, `$CWD/steenrod_modules` and
130///      `ext/steenrod_modules`. The modules can be shifted by appending e.g. `S_2[2]`.
131///    - `module_spec`, a single `&str` of the form `module_name@algebra`, where `module_name` and
132///      `algebra` are as above.
133///  - `save_file`: The save file for the module. If it points to an invalid save file, an error is
134///    returned.
135///
136/// This dispatches to either [`construct_nassau`] or [`construct_standard`] depending on whether
137/// the `nassau` feature is enabled.
138pub fn construct<T, E>(
139    module_spec: T,
140    save_dir: impl Into<SaveDirectory>,
141) -> anyhow::Result<QueryModuleResolution>
142where
143    anyhow::Error: From<E>,
144    T: TryInto<Config, Error = E>,
145{
146    #[cfg(feature = "nassau")]
147    {
148        construct_nassau(module_spec, save_dir)
149    }
150
151    #[cfg(not(feature = "nassau"))]
152    {
153        construct_standard(module_spec, save_dir)
154    }
155}
156
157/// See [`construct`]
158pub fn construct_nassau<T, E>(
159    module_spec: T,
160    save_dir: impl Into<SaveDirectory>,
161) -> anyhow::Result<crate::nassau::Resolution<FDModule<MilnorAlgebra>>>
162where
163    anyhow::Error: From<E>,
164    T: TryInto<Config, Error = E>,
165{
166    let Config {
167        module: json,
168        algebra,
169    } = module_spec.try_into()?;
170
171    if algebra == AlgebraType::Adem {
172        return Err(anyhow!("Nassau's algorithm requires Milnor's basis"));
173    }
174    if !json["profile"].is_null() {
175        return Err(anyhow!(
176            "Nassau's algorithm does not support non-trivial profile"
177        ));
178    }
179    if json["p"].as_i64() != Some(2) {
180        return Err(anyhow!("Nassau's algorithm does not support odd primes"));
181    }
182    if json["type"].as_str() != Some("finite dimensional module") {
183        return Err(anyhow!(
184            "Nassau's algorithm only supports finite dimensional modules"
185        ));
186    }
187
188    let algebra = Arc::new(MilnorAlgebra::new(fp::prime::TWO, false));
189    let module = Arc::new(FDModule::from_json(Arc::clone(&algebra), &json)?);
190
191    if !json["cofiber"].is_null() {
192        return Err(anyhow!("Nassau's algorithm does not support cofiber"));
193    }
194    crate::nassau::Resolution::new_with_save(module, save_dir)
195}
196
197/// See [`construct`]
198pub fn construct_standard<const U: bool, T, E>(
199    module_spec: T,
200    save_dir: impl Into<SaveDirectory>,
201) -> anyhow::Result<crate::resolution::MuResolution<U, CCC>>
202where
203    anyhow::Error: From<E>,
204    T: TryInto<Config, Error = E>,
205    SteenrodAlgebra: algebra::MuAlgebra<U>,
206{
207    let Config {
208        module: json,
209        algebra,
210    } = module_spec.try_into()?;
211
212    let algebra = Arc::new(SteenrodAlgebra::from_json(&json, algebra, U)?);
213    let module = Arc::new(steenrod_module::from_json(Arc::clone(&algebra), &json)?);
214    let mut chain_complex = Arc::new(FiniteChainComplex::ccdz(Arc::clone(&module)));
215
216    let cofiber = &json["cofiber"];
217    if !cofiber.is_null() {
218        assert!(!U, "Cofiber not supported for unstable resolution");
219        use algebra::module::homomorphism::FreeModuleHomomorphism;
220
221        use crate::{chain_complex::ChainMap, yoneda::yoneda_representative};
222
223        let shift = json["shift"].as_i64().unwrap_or(0) as i32;
224
225        let cofiber = BidegreeGenerator::s_t(
226            cofiber["s"].as_i64().unwrap() as i32,
227            cofiber["t"].as_i64().unwrap() as i32 + shift,
228            cofiber["idx"].as_u64().unwrap() as usize,
229        );
230
231        let max_degree = Bidegree::n_s(
232            module
233                .max_degree()
234                .expect("Can only take cofiber when module is bounded"),
235            0,
236        );
237
238        let resolution = Resolution::new(Arc::clone(&chain_complex));
239        resolution.compute_through_stem(cofiber.degree() + max_degree);
240
241        let map = FreeModuleHomomorphism::new(
242            resolution.module(cofiber.s()),
243            Arc::clone(&module),
244            cofiber.t(),
245        );
246        let mut new_output = fp::matrix::Matrix::new(
247            module.prime(),
248            resolution
249                .module(cofiber.s())
250                .number_of_gens_in_degree(cofiber.t()),
251            1,
252        );
253        new_output.row_mut(cofiber.idx()).set_entry(0, 1);
254
255        map.add_generators_from_matrix_rows(cofiber.t(), new_output.as_slice_mut());
256        map.extend_by_zero((max_degree + cofiber.degree()).t());
257
258        let cm = ChainMap {
259            s_shift: cofiber.s(),
260            chain_maps: vec![map],
261        };
262        let yoneda = yoneda_representative(Arc::new(resolution), cm);
263        let mut yoneda = FiniteChainComplex::from(yoneda);
264        yoneda.pop();
265
266        chain_complex = Arc::new(yoneda.map(|m| Box::new(m.clone()) as SteenrodModule));
267    }
268
269    crate::resolution::MuResolution::new_with_save(chain_complex, save_dir)
270}
271
272/// Load a module specification from a JSON file.
273///
274/// Given the name of a module file (without the `.json` extension), find a json file with this
275/// name, and return the parsed json object. The search path for this json file is described
276/// [here](../index.html#module-specification).
277pub fn load_module_json(name: &str) -> anyhow::Result<Value> {
278    let current_dir = std::env::current_dir().context("Failed to read current directory")?;
279    let relative_dir = current_dir.join("steenrod_modules");
280
281    for path in &[
282        current_dir,
283        relative_dir,
284        PathBuf::from(STATIC_MODULES_PATH),
285    ] {
286        let mut path = path.clone();
287        path.push(name);
288        path.set_extension("json");
289        if let Ok(s) = std::fs::read_to_string(&path) {
290            return serde_json::from_str(&s)
291                .with_context(|| format!("Failed to load module json at {path:?}"));
292        }
293    }
294    Err(anyhow!("Module file '{}' not found", name))
295}
296
297/// Given an `n: usize`, return a UTF-8 character that best depicts this number. If `n < 9`, then
298/// this is a UTF-8 when `n` many dots. If `n = 9`, then this is the number `9`. Otherwise, it is
299/// `*`.
300pub fn unicode_num(n: usize) -> char {
301    match n {
302        0 => ' ',
303        1 => '·',
304        2 => ':',
305        3 => '∴',
306        4 => '⁘',
307        5 => '⁙',
308        6 => '⠿',
309        7 => '⡿',
310        8 => '⣿',
311        9 => '9',
312        _ => '*',
313    }
314}
315
316/// Query the user for a module and its save directory. See
317/// [here](../index.html#module-specification) for details on the propmt format.
318///
319/// # Arguments
320/// - `prompt`: The prompt used to query the user for the module. This is `"Module"` when invoked
321///   through [`query_module`], but the user may want to use something more specific, e.g. `"Source
322///   module"`.
323/// - `algebra`: The Steenrod algebra basis allowed. Some applications only support using one of
324///   the two basis, and specifying this parameter forbids the user from specifying the other
325///   basis.
326/// - `load_quasi_inverse`: Whether or not the quasi-inverses of the resolution should be stored.
327///   Note that if there is a save directory, then quasi-inverses will never be stored in memory;
328///   they must be accessed via `apply_quasi_inverse`.
329///
330/// # Returns
331/// A [`QueryModuleResolution`]. Note that this type depends on whether the `nassau` feature is
332/// enabled.
333pub fn query_module_only(
334    prompt: &str,
335    algebra: Option<AlgebraType>,
336    load_quasi_inverse: bool,
337) -> anyhow::Result<QueryModuleResolution> {
338    let (name, module): (String, Config) = query::with_default(prompt, "S_2", |s| {
339        Result::<_, anyhow::Error>::Ok((
340            s.to_owned(),
341            match algebra {
342                Some(algebra) => (s, algebra).try_into()?,
343                None => s.try_into()?,
344            },
345        ))
346    });
347
348    let save_dir = query::optional(&format!("{prompt} save directory"), |x| {
349        core::result::Result::<PathBuf, std::convert::Infallible>::Ok(PathBuf::from(x))
350    });
351
352    let mut resolution =
353        construct(module, save_dir).context("Failed to load module from save file")?;
354
355    let load_quasi_inverse = load_quasi_inverse && resolution.save_dir().is_none();
356
357    #[cfg(not(feature = "nassau"))]
358    {
359        resolution.load_quasi_inverse = load_quasi_inverse;
360    }
361
362    #[cfg(feature = "nassau")]
363    assert!(
364        !load_quasi_inverse,
365        "Quasi inverse loading not support with Nassau. Please use a save directory instead"
366    );
367
368    resolution.set_name(name);
369
370    Ok(resolution)
371}
372
373/// Query the user for a module and a bidegree, and return a resolution resolved up to said
374/// bidegree.
375///
376/// This is mainly a wrapper around [`query_module_only`] that also asks for the bidegree to resolve
377/// up to as well. The prompt of [`query_module_only`] is always set to `"Module"` when invoked
378/// through this function.
379pub fn query_module(
380    algebra: Option<AlgebraType>,
381    load_quasi_inverse: bool,
382) -> anyhow::Result<QueryModuleResolution> {
383    let resolution = query_module_only("Module", algebra, load_quasi_inverse)?;
384
385    let mut max = Bidegree::n_s(
386        query::with_default("Max n", "30", str::parse),
387        query::with_default("Max s", "7", str::parse),
388    );
389
390    if let Some(s) = secondary_job() {
391        if s <= max.s() {
392            max = Bidegree::n_s(max.n(), std::cmp::min(s + 1, max.s()));
393        } else {
394            return Err(anyhow!("SECONDARY_JOB is larger than max_s"));
395        }
396    }
397
398    resolution.compute_through_stem(max);
399
400    Ok(resolution)
401}
402
403pub fn query_unstable_module_only() -> anyhow::Result<SteenrodModule> {
404    let spec: Config = query::raw("Module", |x| x.try_into());
405    let algebra = Arc::new(SteenrodAlgebra::from_json(
406        &spec.module,
407        spec.algebra,
408        true,
409    )?);
410    steenrod_module::from_json(algebra, &spec.module)
411}
412
413pub fn query_unstable_module(load_quasi_inverse: bool) -> anyhow::Result<UnstableResolution<CCC>> {
414    let module = Arc::new(query_unstable_module_only()?);
415    let cc = Arc::new(FiniteChainComplex::ccdz(module));
416
417    let save_dir = query::optional("Module save directory", |x| {
418        core::result::Result::<PathBuf, std::convert::Infallible>::Ok(PathBuf::from(x))
419    });
420
421    let mut resolution = UnstableResolution::new_with_save(cc, save_dir)?;
422    resolution.load_quasi_inverse = load_quasi_inverse && resolution.save_dir().is_none();
423
424    Ok(resolution)
425}
426
427/// Given a resolution, return a resolution of the unit.
428///
429/// The return value comes with a boolean indicating whether the original resolution was already a
430/// resolution of the unit. If the boolean is true, then the original resolution is returned.
431pub fn get_unit(
432    resolution: Arc<QueryModuleResolution>,
433) -> anyhow::Result<(bool, Arc<QueryModuleResolution>)> {
434    let is_unit = resolution.target().max_s() == 1 && resolution.target().module(0).is_unit();
435
436    let unit = if is_unit {
437        Arc::clone(&resolution)
438    } else {
439        let save_dir = query::optional("Unit save directory", |x| {
440            core::result::Result::<PathBuf, std::convert::Infallible>::Ok(PathBuf::from(x))
441        });
442
443        let algebra = resolution.algebra();
444        let module = FDModule::new(
445            algebra,
446            String::from("unit"),
447            bivec::BiVec::from_vec(0, vec![1]),
448        );
449
450        #[cfg(feature = "nassau")]
451        {
452            Arc::new(crate::nassau::Resolution::new_with_save(
453                Arc::new(module),
454                save_dir,
455            )?)
456        }
457
458        #[cfg(not(feature = "nassau"))]
459        {
460            let cc = FiniteChainComplex::ccdz(Arc::new(Box::new(module) as SteenrodModule));
461            Arc::new(Resolution::new_with_save(Arc::new(cc), save_dir)?)
462        }
463    };
464
465    Ok((is_unit, unit))
466}
467
468mod logging {
469    use std::io;
470
471    pub struct LogWriter<T> {
472        writer: T,
473        bytes: u64,
474        start: std::time::Instant,
475    }
476
477    impl<T: io::Write> io::Write for LogWriter<T> {
478        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
479            let written = self.writer.write(buf)?;
480            self.bytes += written as u64;
481            Ok(written)
482        }
483
484        fn flush(&mut self) -> io::Result<()> {
485            self.writer.flush()
486        }
487    }
488
489    impl<T> LogWriter<T> {
490        pub fn new(writer: T) -> Self {
491            Self {
492                writer,
493                bytes: 0,
494                start: std::time::Instant::now(),
495            }
496        }
497    }
498
499    pub struct Throughput(f64);
500
501    impl std::fmt::Display for Throughput {
502        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
503            write!(f, "{:.2} MiB/s", self.0)
504        }
505    }
506
507    impl<T: io::Write> LogWriter<T> {
508        /// Return the throughput in MiB/s
509        pub fn into_throughput(mut self) -> Throughput {
510            self.writer.flush().unwrap();
511            let duration = self.start.elapsed();
512            let mib = self.bytes as f64 / (1024 * 1024) as f64;
513            Throughput(mib / duration.as_secs_f64())
514        }
515    }
516
517    #[cfg(feature = "logging")]
518    pub fn ext_tracing_subscriber() -> impl tracing::Subscriber {
519        use std::io::IsTerminal;
520
521        use tracing_subscriber::{
522            filter::EnvFilter,
523            fmt::{Subscriber, format::FmtSpan},
524        };
525
526        Subscriber::builder()
527            .with_ansi(std::io::stderr().is_terminal())
528            .with_writer(std::io::stderr)
529            .with_max_level(tracing::Level::INFO)
530            .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
531            .with_thread_ids(true)
532            .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_default())
533            .finish()
534    }
535
536    #[cfg(not(feature = "logging"))]
537    pub fn ext_tracing_subscriber() -> impl tracing::Subscriber {
538        tracing::subscriber::NoSubscriber::new()
539    }
540
541    pub fn init_logging() -> anyhow::Result<()> {
542        tracing::subscriber::set_global_default(ext_tracing_subscriber())?;
543
544        tracing::info!("Logging initialized");
545        Ok(())
546    }
547}
548
549pub use logging::{LogWriter, ext_tracing_subscriber, init_logging};
550
551pub(crate) mod parallel {
552
553    use std::sync::atomic::{AtomicUsize, Ordering};
554
555    static PARALLEL_DEPTH: AtomicUsize = AtomicUsize::new(0);
556
557    /// RAII guard that increments [`PARALLEL_DEPTH`] on creation and decrements on drop. Used to mark
558    /// regions where `par_iter_mut` work is active, so that `step_resolution` jobs can detect priority
559    /// inversion and retry.
560    pub(crate) struct ParallelGuard;
561
562    impl ParallelGuard {
563        pub(crate) fn new() -> Self {
564            PARALLEL_DEPTH.fetch_add(1, Ordering::Release);
565            Self
566        }
567    }
568
569    impl Drop for ParallelGuard {
570        fn drop(&mut self) {
571            PARALLEL_DEPTH.fetch_sub(1, Ordering::Release);
572        }
573    }
574
575    pub(crate) fn is_in_parallel() -> bool {
576        PARALLEL_DEPTH.load(Ordering::Acquire) > 0
577    }
578}
579
580/// The value of the SECONDARY_JOB environment variable.
581///
582/// This is used for distributing the `secondary`. If set, only data with `s = SECONDARY_JOB` will
583/// be computed. The minimum value of `s` is the `shift_s` of the
584/// [`SecondaryLift`](crate::secondary::SecondaryLift) and the maximum value (inclusive) is the
585/// maximum `s` of the resolution.
586pub fn secondary_job() -> Option<i32> {
587    let val = std::env::var("SECONDARY_JOB").ok()?;
588    let parsed: Option<i32> = str::parse(&val).ok();
589    if parsed.is_none() {
590        eprintln!(
591            "Invalid argument for `SECONDARY_JOB`. Expected non-negative integer but found {val}"
592        );
593    }
594    parsed
595}