From ba4807e1de44a234fc9499dad0011440cde113fb Mon Sep 17 00:00:00 2001
From: Vincent Ambo <mail@tazj.in>
Date: Wed, 17 May 2023 18:33:22 +0300
Subject: [PATCH] refactor(tvix/cli): split CLI-specific IO logic into TvixIO
 type

This adds a wrapper type TvixIO<T: EvalIO>, which can wrap around an
arbitrary EvalIO implementation and perform actions needed for the
Tvix CLI (marking imported paths as known, and handling __corepkgs__).

Change-Id: I5fc1ca199b9f94b21a89103b84575e0f8f58dff9
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8579
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
---
 tvix/cli/src/main.rs       |  6 ++-
 tvix/cli/src/nix_compat.rs | 11 +-----
 tvix/cli/src/tvix_io.rs    | 78 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 84 insertions(+), 11 deletions(-)
 create mode 100644 tvix/cli/src/tvix_io.rs

diff --git a/tvix/cli/src/main.rs b/tvix/cli/src/main.rs
index 99b5b451f..9a81cff50 100644
--- a/tvix/cli/src/main.rs
+++ b/tvix/cli/src/main.rs
@@ -3,6 +3,7 @@ mod errors;
 mod known_paths;
 mod nix_compat;
 mod refscan;
+mod tvix_io;
 
 use std::cell::RefCell;
 use std::rc::Rc;
@@ -62,7 +63,10 @@ fn interpret(code: &str, path: Option<PathBuf>, args: &Args, explain: bool) -> b
     let known_paths: Rc<RefCell<KnownPaths>> = Default::default();
 
     eval.strict = args.strict;
-    eval.io_handle = Box::new(nix_compat::NixCompatIO::new(known_paths.clone()));
+    eval.io_handle = Box::new(tvix_io::TvixIO::new(
+        known_paths.clone(),
+        nix_compat::NixCompatIO::new(),
+    ));
 
     // bundle fetchurl.nix (used in nixpkgs) by resolving <nix> to
     // `/__corepkgs__`, which has special handling in [`nix_compat`].
diff --git a/tvix/cli/src/nix_compat.rs b/tvix/cli/src/nix_compat.rs
index 1c9f6ba54..70bddb834 100644
--- a/tvix/cli/src/nix_compat.rs
+++ b/tvix/cli/src/nix_compat.rs
@@ -9,10 +9,8 @@ use std::cell::RefCell;
 use std::collections::HashMap;
 use std::path::Path;
 use std::process::Command;
-use std::rc::Rc;
 use std::{io, path::PathBuf};
 
-use crate::known_paths::KnownPaths;
 use smol_str::SmolStr;
 use tvix_eval::{ErrorKind, EvalIO, FileType, StdIO};
 
@@ -23,10 +21,6 @@ pub struct NixCompatIO {
     /// instead.
     underlying: StdIO,
 
-    /// Ingested paths must be reported to this known paths tracker
-    /// for accurate build reference scanning.
-    known_paths: Rc<RefCell<KnownPaths>>,
-
     /// Cache paths for identical files being imported to the store.
     // TODO(tazjin): This could be done better by having a thunk cache
     // for these calls on the eval side, but that is a little more
@@ -86,11 +80,10 @@ impl EvalIO for NixCompatIO {
 }
 
 impl NixCompatIO {
-    pub fn new(known_paths: Rc<RefCell<KnownPaths>>) -> Self {
+    pub fn new() -> Self {
         NixCompatIO {
             underlying: StdIO,
             import_cache: RefCell::new(HashMap::new()),
-            known_paths,
         }
     }
 
@@ -118,8 +111,6 @@ impl NixCompatIO {
             .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
         let out_path_trimmed = out_path_str.trim();
 
-        self.known_paths.borrow_mut().plain(out_path_trimmed);
-
         let mut out_path = PathBuf::new();
         out_path.push(out_path_trimmed);
         Ok(out_path)
diff --git a/tvix/cli/src/tvix_io.rs b/tvix/cli/src/tvix_io.rs
new file mode 100644
index 000000000..401b8a5f4
--- /dev/null
+++ b/tvix/cli/src/tvix_io.rs
@@ -0,0 +1,78 @@
+//! This module implements a wrapper around tvix-eval's [EvalIO] type,
+//! adding functionality which is required by tvix-cli:
+//!
+//! 1. Marking plain paths known to the reference scanner.
+//! 2. Handling the C++ Nix `__corepkgs__`-hack for nixpkgs bootstrapping.
+//!
+//! All uses of [EvalIO] in tvix-cli must make use of this wrapper,
+//! otherwise fundamental features like nixpkgs bootstrapping and hash
+//! calculation will not work.
+
+use crate::KnownPaths;
+use smol_str::SmolStr;
+use std::cell::RefCell;
+use std::path::{Path, PathBuf};
+use std::rc::Rc;
+use tvix_eval::{ErrorKind, EvalIO, FileType};
+
+pub(crate) struct TvixIO<T: EvalIO> {
+    /// Ingested paths must be reported to this known paths tracker
+    /// for accurate build reference scanning.
+    known_paths: Rc<RefCell<KnownPaths>>,
+
+    // Actual underlying [EvalIO] implementation.
+    actual: T,
+}
+
+impl<T: EvalIO> TvixIO<T> {
+    pub(crate) fn new(known_paths: Rc<RefCell<KnownPaths>>, actual: T) -> Self {
+        Self {
+            known_paths,
+            actual,
+        }
+    }
+}
+
+impl<T: EvalIO> EvalIO for TvixIO<T> {
+    fn store_dir(&self) -> Option<String> {
+        self.actual.store_dir()
+    }
+
+    fn import_path(&mut self, path: &Path) -> Result<PathBuf, ErrorKind> {
+        let imported_path = self.actual.import_path(path)?;
+        self.known_paths
+            .borrow_mut()
+            .plain(imported_path.to_string_lossy());
+
+        Ok(imported_path)
+    }
+
+    fn path_exists(&mut self, path: PathBuf) -> Result<bool, ErrorKind> {
+        if path.starts_with("/__corepkgs__") {
+            return Ok(true);
+        }
+
+        self.actual.path_exists(path)
+    }
+
+    fn read_to_string(&mut self, path: PathBuf) -> Result<String, ErrorKind> {
+        // Bundled version of corepkgs/fetchurl.nix. The counterpart
+        // of this happens in `main`, where the `nix_path` of the
+        // evaluation has `nix=/__corepkgs__` added to it.
+        //
+        // This workaround is similar to what cppnix does for passing
+        // the path through.
+        //
+        // TODO: this comparison is bad and allocates, we should use
+        // the sane path library.
+        if path.starts_with("/__corepkgs__/fetchurl.nix") {
+            return Ok(include_str!("fetchurl.nix").to_string());
+        }
+
+        self.actual.read_to_string(path)
+    }
+
+    fn read_dir(&mut self, path: PathBuf) -> Result<Vec<(SmolStr, FileType)>, ErrorKind> {
+        self.actual.read_dir(path)
+    }
+}