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 18) (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)
     (cond
       ((eqv? (store-path-output rt) "")
        (fprintf out "#<store path ~A>" (store-path-path rt)))
       ((drv-is-ca (store-path-drv rt))
        (fprintf out "#<store path ~A (ca~~ ~A!~A)>" (store-path-path rt) (derivation-path (store-path-drv rt)) (store-path-output rt)))
       (else
        (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 (or (drv-is-ca (store-path-drv path)) (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)
       (when (drv-is-ca drv) (error "tried materializing CA 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 '())

     (for-each
       (lambda (drv-and-outputs)
         (define drv (car drv-and-outputs))
         (define ca-drv (drv-resolve-ca (car drv-and-outputs) (cdr drv-and-outputs)))
         (when ca-drv
          (for-each
           (lambda (output)
            (define placeholder (derivation-output-path (cdr (assoc output (derivation-outputs drv)))))
            (define new-path (cdr (assoc output ca-drv)))
            (replace-placeholder placeholder new-path 0))
           (cdr drv-and-outputs))))
       drv-context)

     path)

   (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 (rewrite-bytevector str with-rewrites)
     (define parts '())
     (define (find-part-at i last-i)
       (define next-slash (bytestring-index str (lambda (c) (= c #x2F)) i))
       (if (or (not next-slash) (>= next-slash (- (bytevector-length str) 53)))
         (if (= last-i 0)
           (set! parts #f)
           (set! parts (cons (bytevector-copy str last-i) parts)))
         (let* ((actual-string (utf8->string (bytevector-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 (string->utf8 (cdr mapping-pair)) (cons (bytevector-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)
       (apply bytevector-append (reverse parts))
       str))

   (define (rewrite-string-or-bytevector str with-rewrites)
    (if (bytevector? str)
      (rewrite-bytevector str with-rewrites)
      (rewrite-string str with-rewrites)))

   (define-record-type <pending-item>
     (make-pending-item ca-drv ia-drv resolved-paths awaiting-count awaited-by)
     pending-item?
     (ca-drv pending-item-ca-drv)
     (ia-drv pending-item-ia-drv set-pending-item-ia-drv!)
     (resolved-paths pending-item-resolved-paths set-pending-item-resolved-paths!)
     (awaiting-count pending-item-awaiting-count set-pending-item-awaiting-count!)
     (awaited-by pending-item-awaited-by set-pending-item-awaited-by!))

   (define-record-printer (<pending-item> item out)
    (fprintf out "#<pending-item ~A - awaiting ~S>" (derivation-path (pending-item-ca-drv item)) (pending-item-awaiting-count item)))

   (define (rewrite-ca-stack input-drv)
     ; A mapping of CA derivation path to <pending-item>.
     (define ca-to-pending-map (mapping (make-default-comparator)))

     (define pending-mutex (make-mutex))
     (define pending-count 0)
     (define build-mutex (make-mutex))
     (define build-condvar (make-condition-variable))
     (define to-build '())
     (define (depend-on pend item)
       (set-pending-item-awaiting-count! pend (+ (pending-item-awaiting-count pend) 1))
       (set-pending-item-awaited-by! item (cons pend (pending-item-awaited-by item))))

     (define (handle-new-drv pend)
      (for-each
       (lambda (input-kv)
        (when (drv-is-ca (car input-kv))
         (depend-on pend (get-item (car input-kv)))))
       (derivation-input-drvs (pending-item-ca-drv pend)))
      (wake-up pend)
      pend)

     (define (wake-up pend)
      (when (eq? (pending-item-awaiting-count pend) 0)
       (set-pending-item-awaiting-count! pend 'build)
       (mutex-lock! build-mutex)
       (set! to-build (cons pend to-build))
       (condition-variable-signal! build-condvar)
       (mutex-unlock! build-mutex)))

     (define (get-item drv)
      (unless (drv-is-ca drv) (error "drv not CA" (derivation-path drv)))
      (define pending (mapping-ref/default ca-to-pending-map (derivation-path drv) #f))
      (unless pending
        (set! pending (make-pending-item drv #f '() 0 '()))
        (mutex-lock! pending-mutex)
        (set! pending-count (+ pending-count 1))
        (set! ca-to-pending-map (mapping-set! ca-to-pending-map (derivation-path drv) pending))
        (mutex-unlock! pending-mutex)
        (handle-new-drv pending))
      pending)

     (define (do-build conn item)
       ; Rewrite CA drv to IA drv using the known inputs
       (define new-drvs (list))
       (define new-srcs (derivation-input-src (pending-item-ca-drv item)))
       (define rewrites (list))
       (for-each
        (lambda (drv-and-outputs)
          (mutex-lock! pending-mutex)
          (define dep-pend (mapping-ref/default ca-to-pending-map (derivation-path (car drv-and-outputs)) #f))
          (mutex-unlock! pending-mutex)
          (if dep-pend
            ; Iterate over each output path, and add its CA equivalent to the input list here.
            (for-each
             (lambda (output)
              (define new-path (cdr (assoc output (pending-item-resolved-paths dep-pend))))
              (set! new-srcs (cons new-path new-srcs))
              (define old-output (cdr (assoc output (derivation-outputs (car drv-and-outputs)))))
              (set! rewrites (cons (cons (derivation-output-path old-output) new-path) rewrites)))
             (cdr drv-and-outputs))

            ; Not a CA drv, add it back to the drvs list
            (set! new-drvs (cons drv-and-outputs new-drvs))))
        (derivation-input-drvs (pending-item-ca-drv item)))

       (define ca-drv (pending-item-ca-drv item))
       (define new-builder (rewrite-string-or-bytevector (derivation-builder ca-drv) rewrites))
       (define new-args (map (lambda (v) (rewrite-string-or-bytevector v rewrites)) (derivation-args ca-drv)))
       (define new-env (map (lambda (kv) (cons (car kv) (rewrite-string-or-bytevector (cdr kv) rewrites))) (derivation-env ca-drv)))
       (define ia-drv
        (make-input-addressed-derivation
          (derivation-name ca-drv)
          (derivation-system ca-drv)
          (list-sort
            (lambda (l r) (string<? (derivation-path (car l)) (derivation-path (car r))))
            new-drvs)
          (list-sort string<? new-srcs)
          (cons new-builder new-args) new-env (map car (derivation-outputs ca-drv))))
       (set-pending-item-ia-drv! item ia-drv)

       ; Build all the paths.
       (parameterize ((*daemon* conn)) (write-drv-to-daemon ia-drv))
       (define outputs (map car (derivation-outputs ia-drv)))
       (daemon-wop-build-paths conn (list->vector (map (lambda (v) (string-append (derivation-path ia-drv) "!" v)) outputs)))
       (set-pending-item-resolved-paths! item
        (map (lambda (o)
               (define-values (name hash nar-size ca-store-path)
                 (store-path-to-fod conn (derivation-output-path (cdr (assoc o (derivation-outputs ia-drv) string=?)))))
               (cons o ca-store-path))
         outputs))

       ; Notify our dependencies that we're done.
       (mutex-lock! pending-mutex)
       (set! pending-count (- pending-count 1))
       (mutex-unlock! pending-mutex)
       (set-pending-item-awaiting-count! item 'built)
       (for-each (lambda (depends-on) (set-pending-item-awaiting-count! depends-on (- (pending-item-awaiting-count depends-on) 1)) (wake-up depends-on)) (pending-item-awaited-by item)))

     (define root-pend (get-item input-drv))
     (define (builder conn)
       (mutex-lock! build-mutex)
       (define item #f)
       (when (pair? to-build)
         (set! item (car to-build))
         (set! to-build (cdr to-build)))
       (define local-pending-count #f)
       (if item
         ; If we got an item: unlock the build mutex and build it
         (begin (mutex-unlock! build-mutex) (do-build conn item) (builder conn))
         (begin
           ; Check how many pending items there are..
           (mutex-lock! pending-mutex)
           (set! local-pending-count pending-count)
           (mutex-unlock! pending-mutex)
           (if (= local-pending-count 0)
             ; We're out of pending items, unlock the mutex and drop the thread
             (begin
              (mutex-unlock! build-mutex)
              ; Notify the other threads that we have no more builds to do.
              (condition-variable-broadcast! build-condvar))

             ; We still have pending items, let's go back and wait.
             (begin
              (mutex-unlock! build-mutex build-condvar)
              (builder conn))))))

     (define builder-threads (list))
     (do ((i 0 (+ i 1))) ((>= i 16) #f)
       (set! builder-threads (cons (thread-start! (make-thread (lambda () (builder (daemon-connect))) (string-append "ca-builder-" (number->string i)))) builder-threads)))

     (for-each thread-join! builder-threads)
     root-pend)
   (define (drv-resolve-ca drv outputs)
     (if (drv-is-ca drv)
       (pending-item-resolved-paths (rewrite-ca-stack 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)
       (set! val (resolve-upstream-output-placeholders val (zexp-evaluation-drvs ctx)))
       (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)))))