zilch/core/src/magic.sld

;; Defines procedures to interact with the Nix store by way of zexpressions.
;; This library defines the `<store-path>` record type, which can be used in zexps.
;; A `<store-path>` unquotes in `zexp`s as its store path.
(define-library (zilch magic)
  (import
    (scheme base) (scheme file)
    (zilch lib hash) (zilch nix daemon) (zilch nix drv) (zilch nix path)
    (zilch nix hash)
    (zilch planner step)
    (zilch zexpr)
    (srfi 128) (srfi 132) (srfi 146) (srfi 152) (srfi 207)
    (chicken base) (chicken format) socket)

  (export
    *daemon* *use-ca*
    <store-path>
    make-store-path store-path?
    store-path-drv store-path-output

    store-path-path store-path-build store-path-materialize store-path-realisation
    store-path-for-text store-path-for-fod store-path-for-drv
    store-path-for-impure-drv store-path-for-ca-drv store-path-for-ca-drv*
    store-path-realised store-path-open

    drv-resolve-ca

    zilch-magic-counters)

  (begin

   (define (daemon-connect)
     (define conn
      (parameterize
       ((socket-send-buffer-size 4096) (socket-send-size 4096) (socket-receive-timeout 60000) (socket-send-timeout 5000))
       (let ((unix-socket (socket af/unix sock/stream)))
            (socket-connect unix-socket (unix-address "/nix/var/nix/daemon-socket/socket"))
            (let-values (((in-port out-port) (socket-i/o-ports unix-socket)))
                        (make-daemon-link in-port out-port)))))
     (daemon-wop-handshake conn)
     conn)

   (define (daemon-close conn)
     (close-input-port (daemon-link-in-port conn))
     (close-output-port (daemon-link-out-port conn)))

   ;; The daemon connection used by `(zilch magic)`.
   (define *daemon*
     (make-parameter (daemon-connect)))

   ; Create a CA store path from the store path being passed in.
   (define (store-path-to-fod conn path)
     (define data (daemon-wop-query-path-info conn path))
     (define nar-size (valid-path-info-nar-size data))
     (define hash (valid-path-info-nar-hash data))
     (define references (valid-path-info-references data))
     (define references-filtered (list-copy references))
     (define self-references (member path references-filtered string=?))
     (when self-references
       (if (null? (cdr self-references))
         (set! references-filtered '())
         (begin
           ; Cheaply remove this item from the list.
           ;    (self-references . (foo . (bar . baz)))
           ; -> (foo             . (bar . baz)
           (set-car! self-references (cadr self-references))
           (set-cdr! self-references (cddr self-references)))))

     (define name (string-copy path (+ (string-length (%store-dir)) 1 32 1)))
     (define ca-store-path (make-fixed-output-with-references hash name references-filtered self-references))
     (unless (daemon-wop-query-path-info conn ca-store-path)
       (daemon-wop-add-to-store-nar conn ca-store-path (valid-path-info-deriver data) (hex hash) references nar-size (string-append "fixed:r:sha256:" (as-base32 hash))
         (lambda (write-blob)
           (define new-conn (daemon-connect))
           (daemon-wop-nar-from-path new-conn path)
           (define blob (make-bytevector 4096))
           (do ((i 0))
               ((= i nar-size) (daemon-close new-conn))
               (let* ((chunk-size (min 4096 (- nar-size i)))
                      (bytes-read (read-bytevector! blob (daemon-link-in-port new-conn) 0 chunk-size)))
                 (when (eof-object? bytes-read)
                   (error "unexpected EOF"))
                 (set! i (+ i bytes-read))
                 (write-blob blob 0 bytes-read))))))
     (values name hash nar-size ca-store-path))

   ;; If set to `#f`, `store-path-for-ca-drv*` will not generate
   ;; content-addressed derivations.
   (define *use-ca* (make-parameter #t))

   ;; A vector of counters, counting the amount of derivations made, built, and IFD'd.
   (define zilch-magic-counters (vector 0 0 0))

   (define (increment-counter index)
     (vector-set! zilch-magic-counters index (+ 1 (vector-ref zilch-magic-counters index))))

   ;; Represents a reference to an output path of a derivation, or a source file.
   ;; if `output` is `""`, `drv` is the store path to a source file.
   (define-record-type <store-path>
     (make-store-path drv output written)
     store-path?
     (drv store-path-drv)
     (output store-path-output)
     (written store-path-written set-store-path-written!))

   (define-record-printer (<store-path> rt out)
     (if (eqv? (store-path-output rt) "")
       (fprintf out "#<store path ~A>" (store-path-path rt))
       (fprintf out "#<store path ~A (~A!~A)>" (store-path-path rt) (derivation-path (store-path-drv rt)) (store-path-output rt))))

   ;; Returns the store path for the output associated with this `<store-path>`.
   (define (store-path-path path)
     (derivation-output-path (cdr (assoc (store-path-output path) (derivation-outputs (store-path-drv path))))))

   ;; Makes sure the derivation referenced by this store path exists in the daemon.
   (define (store-path-materialize path)
     (unless (store-path-written path)
       (write-drv-to-daemon (store-path-drv path))
       (set-store-path-written! path #t)))

   ;; Returns the output path of this store path; fetching it from the daemon if
   ;; the derivation is content-addressed.
   (define (store-path-realisation path)
     (define drv (store-path-drv path))
     (define output (store-path-output path))
     (define drv-output (cdr (assoc output (derivation-outputs drv))))
     (if (or (not (derivation-output-hash drv-output)) (bytevector? (derivation-output-hash drv-output)))
       (derivation-output-path drv-output)
       (begin
         (store-path-materialize path)
         (let ((outputs (daemon-wop-query-derivation-output-map (*daemon*) (derivation-path drv))))
           (cdr (assoc output outputs))))))

   ;; Requests that the daemon build this store path.
   (define (store-path-build path)
     (increment-counter 1)
     (daemon-wop-build-paths (*daemon*) (vector (string-append (derivation-path (store-path-drv path)) "!" (store-path-output path)))))

   ;; Writes the `<derivation>` to the Nix store, via the currently specified `*daemon*`.
   (define (write-drv-to-daemon drv)
       (define path (derivation-path drv))
       (unless (file-exists? path)
        (let ((out (open-output-string)))
            (derivation-serialize drv out)
            (daemon-wop-add-text-to-store (*daemon*) (string-append (derivation-name drv) ".drv") (get-output-string out) (derivation-path-references drv))))
       (make-store-path path "" #t))

   ;; Returns a store path representing the text.
   (define (store-path-for-text name text)
     (increment-counter 0)
     (define goal-path (make-text-path "sha256" (sha256 text) name '()))
     (unless (file-exists? goal-path) (daemon-wop-add-text-to-store (*daemon*) name text '()))
     (make-store-path goal-path "" #t))

   ;; Returns a `<store-path>` for a fixed output derivation.
   (define (store-path-for-fod name platform builder env hash-algo hash-value hash-recursive)
     (increment-counter 0)
     (define collected-env (zexp-unwrap env))
     (define collected-builder (zexp-unwrap builder))

     (define input-drvs (merge-drvs (zexp-evaluation-drvs collected-env) (zexp-evaluation-drvs collected-builder)))
     (define input-srcs (merge-srcs (zexp-evaluation-srcs collected-env) (zexp-evaluation-srcs collected-builder)))
     (define drv (make-fixed-output-derivation name platform input-drvs input-srcs (zexp-evaluation-value collected-builder) (zexp-evaluation-value collected-env) hash-algo hash-value hash-recursive))
     (make-store-path drv "out" #f))

   ;; Returns an alist of output -> `<store-path>` for an input-addressed derivation.
   (define (store-path-for-drv name platform builder env outputs)
     (increment-counter 0)
     (define collected-env (zexp-unwrap env))
     (define collected-builder (zexp-unwrap builder))

     (define input-drvs (merge-drvs (zexp-evaluation-drvs collected-env) (zexp-evaluation-drvs collected-builder)))
     (define input-srcs (merge-srcs (zexp-evaluation-srcs collected-env) (zexp-evaluation-srcs collected-builder)))
     (define drv (make-input-addressed-derivation name platform input-drvs input-srcs (zexp-evaluation-value collected-builder) (zexp-evaluation-value collected-env) outputs))
     (map (lambda (l) (cons (car l) (make-store-path drv (car l) #f))) (derivation-outputs drv)))

   ;; Returns an alist of output -> `<store-path>` for an impure derivation.
   (define (store-path-for-impure-drv name platform builder env outputs)
     (increment-counter 0)
     (define collected-env (zexp-unwrap env))
     (define collected-builder (zexp-unwrap builder))

     (define input-drvs (merge-drvs (zexp-evaluation-drvs collected-env) (zexp-evaluation-drvs collected-builder)))
     (define input-srcs (merge-srcs (zexp-evaluation-srcs collected-env) (zexp-evaluation-srcs collected-builder)))
     (define drv (make-impure-derivation name platform input-drvs input-srcs (zexp-evaluation-value collected-builder) (zexp-evaluation-value collected-env) outputs))
     (map (lambda (l) (cons (car l) (make-store-path drv (car l) #f))) (derivation-outputs drv)))

   ;; Returns an alist of output -> `<store-path>` for a content-addressed derivation.
   (define (store-path-for-ca-drv name platform builder env outputs)
     (increment-counter 0)
     (define collected-env (zexp-unwrap env))
     (define collected-builder (zexp-unwrap builder))

     (define input-drvs (merge-drvs (zexp-evaluation-drvs collected-env) (zexp-evaluation-drvs collected-builder)))
     (define input-srcs (merge-srcs (zexp-evaluation-srcs collected-env) (zexp-evaluation-srcs collected-builder)))
     (define drv (make-ca-derivation name platform input-drvs input-srcs (zexp-evaluation-value collected-builder) (zexp-evaluation-value collected-env) outputs))
     (map (lambda (l) (cons (car l) (make-store-path drv (car l) #f))) (derivation-outputs drv)))

   ;; Calls either `store-path-for-ca-drv` or `store-path-for-drv` depending on `*use-ca*`.
   (define (store-path-for-ca-drv* name platform builder env outputs)
     (if (*use-ca*) (store-path-for-ca-drv name platform builder env outputs)
                    (store-path-for-drv name platform builder env outputs)))

   (define (merge-drvs left right)
     ; Create a new pair for the head of each drvs list
     (define drvs (map (lambda (l) (cons (car l) (cdr l))) left))

     (for-each
       (lambda (item)
         (define left (assoc (car item) drvs derivation-equal?))
         (if left
           (for-each
             (lambda (output)
               (unless (member output (cdr left))
                 (set-cdr! left (cons output (cdr left)))))
             (cdr item))
           (set! drvs (cons item drvs))))
       right)
     (list-sort (lambda (l r) (string<? (derivation-path (car l)) (derivation-path (car r)))) (map (lambda (a) (cons (car a) (list-sort string<? (cdr a)))) drvs)))

   (define (merge-srcs left right)
     (for-each (lambda (item) (when (eq? (member item left) #f) (set! left (cons item left)))) right)
     (list-sort string<? left))

   (define (resolve-upstream-output-placeholders path drv-context)
     (define known-placeholders (mapping (make-default-comparator)))
     ; Returns #t if placeholder was replaced
     (define (replace-placeholder placeholder replacement start-index)
       (define index (string-contains path placeholder start-index))
       (if index
         (begin
           (set! path (string-replace path replacement index (+ index (string-length placeholder))))
           (replace-placeholder placeholder replacement (+ index (string-length replacement))))
         (> start-index 0)))

     (define context-to-build '())
     (define drv-output-map #f)
     (define placeholders-to-build '())

     (define (check-output-needs-building output-name output drv)
       (when (and (derivation-output-placeholder? output) (string-contains path (derivation-output-path output)))
         (unless drv-output-map
           (set! drv-output-map (daemon-wop-query-derivation-output-map (*daemon*) (derivation-path drv))))
         (let ((known-path (cdr (assoc output-name drv-output-map))))
           (unless known-path
             (set! placeholders-to-build (cons (string-append (derivation-path drv) "!" output-name) placeholders-to-build))))))

     (define (process-output output-name output drv)
       (define needs-building #f)
       (when (and (derivation-output-placeholder? output) (string-contains path (derivation-output-path output)))
         (unless drv-output-map
           (set! drv-output-map (daemon-wop-query-derivation-output-map (*daemon*) (derivation-path drv))))
         (let* ((known-path (cdr (assoc output-name drv-output-map)))
                (is-replaced (replace-placeholder (derivation-output-path output) known-path 0)))
           (when (and is-replaced (not (file-exists? known-path)))
             (set! needs-building #t))))
       (when (or needs-building (not (or (derivation-output-placeholder? output) (file-exists? (derivation-output-path output)))))
         (set! context-to-build (cons (string-append (derivation-path drv) "!" output-name) context-to-build))))

     (for-each
       (lambda (drv-outputs)
         (set! drv-output-map #f)
         (for-each
           (lambda (output-name)
             (check-output-needs-building output-name (cdr (assoc output-name (derivation-outputs (car drv-outputs)))) (car drv-outputs)))
           (cdr drv-outputs)))
       drv-context)
     (unless (null? placeholders-to-build)
       (daemon-wop-build-paths (*daemon*) (list->vector placeholders-to-build)))
     (for-each
       (lambda (drv-outputs)
         (set! drv-output-map #f)
         (for-each
           (lambda (output-name)
             (process-output output-name (cdr (assoc output-name (derivation-outputs (car drv-outputs)))) (car drv-outputs)))
           (cdr drv-outputs)))
       drv-context)
     (values path context-to-build))

   (define (zexp-ctx-has-placeholder drv-context)
     (if (null? drv-context)
       #f
       (let ((drv (caar drv-context))
             (outputs (cdar drv-context))
             (has-placeholder #f))
         (for-each (lambda (output) (set! has-placeholder (or has-placeholder (derivation-output-placeholder? (cdr (assoc output (derivation-outputs drv))))))) outputs)
         (or has-placeholder (zexp-ctx-has-placeholder (cdr drv-context))))))

   (define (drv-is-ca drv)
     (define is-ca #f)
     (for-each (lambda (out) (when (eq? (derivation-output-hash (cdr out)) 'floating) (set! is-ca #t))) (derivation-outputs drv))
     is-ca)

   (define (rewrite-string str with-rewrites)
     (define parts '())
     (define (find-part-at i last-i)
       (define next-slash (string-contains str "/" i))
       (if (or (not next-slash) (>= next-slash (- (string-length str) 53)))
         (if (= last-i 0)
           (set! parts #f)
           (set! parts (cons (string-copy str last-i) parts)))
         (let* ((actual-string (string-copy str next-slash (+ next-slash 53)))
                (mapping-pair (assoc actual-string with-rewrites string=?)))
           ; If we have a mapping for this string, replace it and continue.
           (if mapping-pair
             (begin
               (set! parts (cons (cdr mapping-pair) (cons (string-copy str last-i next-slash) parts)))
               (find-part-at (+ next-slash 53) (+ next-slash 53)))
             (find-part-at (+ next-slash 1) last-i)))))
     (find-part-at 0 0)
     (if (pair? parts)
       (string-concatenate-reverse parts)
       str))

   (define (drv-resolve-ca drv needed-outputs)
     (if (drv-is-ca drv)
       (let
         ((new-drvs (list)) (new-srcs (derivation-input-src drv)) (rewrites (list)))
         ; For each drv, figure out if we need to CA it.
         (begin
          (for-each
            (lambda (drv-and-outputs)
              (define ca-outputs (drv-resolve-ca (car drv-and-outputs) (cdr drv-and-outputs)))
              (if (not ca-outputs)
                ; It's not CA'd, so add it back
                (set! new-drvs (cons drv-and-outputs new-drvs))

                ; Find the CA paths matching the input paths (this is basically equivalent to daemon-wop-query-derivation-output-map!)
                ; Also record the rewrites!
                (for-each
                  (lambda (output)
                    (set! new-srcs (cons (cdr (assoc output (cdr ca-outputs))) new-srcs))
                    (define old-output (cdr (assoc output (derivation-outputs (car drv-and-outputs)))))
                    (set! rewrites (cons (cons (derivation-output-path old-output) (cdr (assoc output (cdr ca-outputs)))) rewrites)))
                  (cdr drv-and-outputs))))
           (derivation-input-drvs drv))
          ; Now we rewrite the builder and environment.
          (let*
            ((new-builder (rewrite-string (derivation-builder drv) rewrites))
             (new-args (map (lambda (v) (rewrite-string v rewrites)) (derivation-args drv)))
             (new-env (map (lambda (kv) (cons (car kv) (rewrite-string (cdr kv) rewrites))) (derivation-env drv)))
             (new-drv (make-input-addressed-derivation
                       (derivation-name drv)
                       (derivation-system drv)
                       new-drvs new-srcs (cons new-builder new-args) new-env (map car (derivation-outputs drv)))))
            (write-drv-to-daemon new-drv)
            (daemon-wop-build-paths (*daemon*) (list->vector (map (lambda (v) (string-append (derivation-path new-drv) "!" v)) needed-outputs)))
            (map (lambda (o) (define-values (name hash nar-size ca-store-path) (store-path-to-fod (*daemon*) (derivation-output-path (cdr (assoc o (derivation-outputs new-drv) string=?))))) ca-store-path)
              needed-outputs))))
       ; Not a CA drv.
       #f))

   (define (store-path-realised path)
     (define ctx (zexp-unwrap (zexp (zexp-unquote path))))
     (define val (zexp-evaluation-value ctx))
     (define to-build (list))
     (for-each
      (lambda (drv-and-outputs)
       (unless (drv-is-ca (car drv-and-outputs))
         (for-each
          (lambda (o)
           (set! to-build (cons (string-append (derivation-path (car drv-and-outputs)) "!" o) to-build)))
          (cdr drv-and-outputs))))
      (zexp-evaluation-drvs ctx))
     (if (string? val)
       (let-values (((resolved resolved-to-build) (resolve-upstream-output-placeholders val (zexp-evaluation-drvs ctx))))
         (set! val resolved)
         (set! to-build resolved-to-build))
       (when (zexp-ctx-has-placeholder (zexp-evaluation-drvs ctx))
         (error "store-path-realised: expression has dependencies on placeholder context, but isn't a string" (list path val))))
     (when (and (string? val) (not (file-exists? val)) (not (null? to-build)))
       (daemon-wop-build-paths (*daemon*) (list->vector to-build)))
     val)

   (register-build-step '((zilch core magic) build) #t
    (lambda items
      (printf "received build info: ~S\n" items)
      (define all-paths '())
      (for-each
        (lambda (item)
          (define outputs (list-ref item 1))
          (define drv-path (list-ref item 2))
          (set! all-paths (append (map (lambda (o) (string-append drv-path "!" o)) outputs) all-paths)))
        items)
      (daemon-wop-build-paths (*daemon*) (list->vector all-paths))
      (define output '())
      (for-each
        (lambda (item)
          (define output-map (daemon-wop-query-derivation-output-map (*daemon*) (list-ref item 2)))
          (set! output (cons (cons (car item) output-map) output)))
        items)
      output))

   ;; Ensures the `<store-path>` exists, then opens an input port to allow reading from it.
   (define (store-path-open path)
     (increment-counter 2)
     (define output-path (store-path-realised path))
     (open-input-file output-path))

   (zexp-add-unquote-handler
     (lambda (val)
       (if (store-path? val)
           (begin
             (if (string=? (store-path-output val) "")
                 (begin (zexp-context-register-items '() (list (store-path-drv val))) (store-path-drv val))
                 (begin (store-path-materialize val) (zexp-context-register-items `((,(store-path-drv val) ,(store-path-output val))) '()) (store-path-path val))))
           #f)))))