cs: inline struct accessors/mutators across compilation units

Make schemify inline structure accessors and mutators across linklet
boundaries --- or, in JIT mode, across function boundaries --- by
replacing an accessor or mutator with a `#%record?` test and
`unsafe-struct*-{ref,set!}` operation.

racket/src/cs/primitive/internal.ss

@@ -28,5 +28,6 @@
   [record-predicate (known-constant)]
   [record-accessor (known-constant)]
   [record-mutator (known-constant)]
+  [unsafe-struct? (known-constant)]
 
   [make-pthread-parameter (known-procedure 2)])

racket/src/cs/rumble.sls

@@ -619,6 +619,7 @@
           unsafe-struct-set!
           unsafe-struct*-ref
           unsafe-struct*-set!
+          unsafe-struct? ; not exported to racket
 
           unsafe-s16vector-ref
           unsafe-s16vector-set!

racket/src/cs/rumble/struct.ss

@@ -925,6 +925,8 @@
   (#3%vector-ref s i))
 (define (unsafe-struct*-set! s i v)
   (#3%vector-set! s i v))
+(define (unsafe-struct? v r)
+  (#3%record? v r))
 
 (define (unsafe-struct-ref s i)
   (if (impersonator? s)

racket/src/schemify/find-definition.rkt

@@ -45,13 +45,26 @@
                                 (unwrap s?)
                                 (known-predicate 2 type))]
               [knowns
+               (let* ([immediate-count (struct-type-info-immediate-field-count info)]
+                      [parent-count (- (struct-type-info-field-count info)
+                                       immediate-count)])
                  (for/fold ([knowns knowns]) ([id (in-list acc/muts)]
                                               [maker (in-list make-acc/muts)])
+                   (match maker
+                     [`(,make ,ref-or-set ,pos (quote ,name))
+                      (or (and (exact-nonnegative-integer? pos)
+                               (pos . < . immediate-count)
+                               (symbol? name)
                                (cond
-                  [(wrap-eq? (wrap-car maker) -ref)
+                                 [(and (wrap-eq? make 'make-struct-field-accessor)
-                   (hash-set knowns (unwrap id) (known-accessor 2 type))]
+                                       (wrap-eq? ref-or-set -ref))
-                  [else
+                                  (hash-set knowns (unwrap id) (known-field-accessor 2 type struct:s (+ parent-count pos)))]
-                   (hash-set knowns (unwrap id) (known-mutator 4 type))]))])
+                                 [(and (wrap-eq? make 'make-struct-field-mutator)
+                                       (wrap-eq? ref-or-set -set!))
+                                  (hash-set knowns (unwrap id) (known-field-mutator 4 type struct:s (+ parent-count pos)))]
+                                 [else knowns]))
+                          knowns)]
+                     [`,_ knowns])))])
          (values (hash-set knowns (unwrap struct:s) (known-struct-type type
                                                                        (struct-type-info-field-count info)
                                                                        (struct-type-info-pure-constructor? info)))

racket/src/schemify/inline.rkt

@@ -8,6 +8,7 @@
 (provide init-inline-fuel
          can-inline?
          inline-clone
+         inline-type-id
          known-inline->export-known)
 
 (define inline-base 3)
@@ -73,6 +74,33 @@
                    (inline-clone i-k im add-import! mutated imports)))]
         [else #f])])))
 
+(define (inline-type-id k im add-import! mutated imports)
+  (define type-id (cond
+                    [(known-field-accessor? k)
+                     (known-field-accessor-type-id k)]
+                    [(known-field-mutator? k)
+                     (known-field-mutator-type-id k)]
+                    [else #f]))
+  (define env
+    ;; A `needed->env` setup can fail if a needed import cannot be
+    ;; made available:
+    (cond
+      [(not type-id) #f]
+      [(not im) '()]
+      [(known-field-accessor/need-imports? k)
+       (needed->env (known-field-accessor/need-imports-needed k)
+                    add-import!
+                    im)]
+      [(known-field-mutator/need-imports? k)
+       (needed->env (known-field-mutator/need-imports-needed k)
+                    add-import!
+                    im)]
+      [else '()]))
+  (and env
+       (cond
+         [(null? env) type-id]
+         [else (clone-expr type-id env mutated)])))
+
 ;; Build a mapping from ids in the expr to imports into the current
 ;; linklet, where `add-import!` arranges for the import to exist as
 ;; needed and if possible. The result is #f if some import cannot be
@@ -187,8 +215,31 @@
         (known-procedure/can-inline/need-imports
          (known-procedure-arity-mask k)
          (known-procedure/can-inline-expr k)
-         (for/list ([(k v) (in-hash needed)])
+         (hash->list needed))])]
-           (cons k v)))])]
+    [(known-field-accessor? k)
+     (define needed (needed-imports (known-field-accessor-type-id k) prim-knowns imports exports '() '#hasheq()))
+     (cond
+       [needed
+        (known-field-accessor/need-imports (known-procedure-arity-mask k)
+                                           (known-accessor-type k)
+                                           (known-field-accessor-type-id k)
+                                           (known-field-accessor-pos k)
+                                           (hash->list needed))]
+       [else
+        (known-accessor (known-procedure-arity-mask k)
+                        (known-accessor-type k))])]
+    [(known-field-mutator? k)
+     (define needed (needed-imports (known-field-mutator-type-id k) prim-knowns imports exports '() '#hasheq()))
+     (cond
+       [needed
+        (known-field-mutator/need-imports (known-procedure-arity-mask k)
+                                          (known-mutator-type k)
+                                          (known-field-mutator-type-id k)
+                                          (known-field-mutator-pos k)
+                                          (hash->list needed))]
+       [else
+        (known-mutator (known-procedure-arity-mask k)
+                       (known-mutator-type k))])]
     [else k]))
 
 (define (needed-imports v prim-knowns imports exports env needed)
@@ -277,3 +328,7 @@
                (wrap-cdr args))]
     [else
      (cons (unwrap args) env)]))
+
+(define (hash->list needed)
+  (for/list ([(k v) (in-hash needed)])
+    (cons k v)))

racket/src/schemify/known.rkt

@@ -18,6 +18,10 @@
          known-predicate known-predicate? known-predicate-type
          known-accessor known-accessor? known-accessor-type
          known-mutator known-mutator? known-mutator-type
+         known-field-accessor known-field-accessor? known-field-accessor-type-id known-field-accessor-pos
+         known-field-mutator known-field-mutator? known-field-mutator-type-id known-field-mutator-pos
+         known-field-accessor/need-imports known-field-accessor/need-imports? known-field-accessor/need-imports-needed
+         known-field-mutator/need-imports known-field-mutator/need-imports? known-field-mutator/need-imports-needed
          known-struct-type-property/immediate-guard known-struct-type-property/immediate-guard?
          a-known-constant
          a-known-consistent)
@@ -56,6 +60,10 @@
 (struct known-predicate (type) #:prefab #:omit-define-syntaxes #:super struct:known-procedure/succeeds)
 (struct known-accessor (type) #:prefab #:omit-define-syntaxes #:super struct:known-procedure)
 (struct known-mutator (type) #:prefab #:omit-define-syntaxes #:super struct:known-procedure)
+(struct known-field-accessor (type-id pos) #:prefab #:omit-define-syntaxes #:super struct:known-accessor)
+(struct known-field-mutator (type-id pos) #:prefab #:omit-define-syntaxes #:super struct:known-mutator)
+(struct known-field-accessor/need-imports (needed) #:prefab #:omit-define-syntaxes #:super struct:known-field-accessor)
+(struct known-field-mutator/need-imports (needed) #:prefab #:omit-define-syntaxes #:super struct:known-field-mutator)
 
 (struct known-struct-type-property/immediate-guard () #:prefab #:omit-define-syntaxes)
 

racket/src/schemify/schemify.rkt

@@ -204,7 +204,7 @@
                                      prim-knowns knowns mutated imports exports
                                      allow-set!-undefined?
                                      add-import!
-                                     for-cify?
+                                     for-cify? for-jitify?
                                      unsafe-mode?))
         (match form
           [`(define-values ,ids ,_)
@@ -259,7 +259,7 @@
 ;; Schemify `let-values` to `let`, etc., and
 ;; reorganize struct bindings.
 (define (schemify v prim-knowns knowns mutated imports exports allow-set!-undefined? add-import!
-                  for-cify? unsafe-mode?)
+                  for-cify? for-jitify? unsafe-mode?)
   (let schemify/knowns ([knowns knowns] [inline-fuel init-inline-fuel] [v v])
     (let schemify ([v v])
       (define s-v
@@ -278,7 +278,7 @@
                          ,make2
                          ,?2
                          ,make-acc/muts ...)))
-            #:guard (not for-cify?)
+            #:guard (not (or for-jitify? for-cify?))
             ;; Convert a `make-struct-type` binding into a 
             ;; set of bindings that Chez's cp0 recognizes,
             ;; and push the struct-specific extra work into
@@ -566,19 +566,74 @@
                           (left-left-lambda-convert
                            (inline-clone k (hash-ref imports u-rator #f) add-import! mutated imports)
                            (sub1 inline-fuel))))))
+            (define (maybe-tmp e name)
+              ;; use `e` directly if it's ok to duplicate
+              (if (simple/can-copy? e prim-knowns knowns imports mutated)
+                  e
+                  (gensym name)))
+            (define (wrap-tmp tmp e body)
+              (if (eq? tmp e)
+                  body
+                  `(let ([,tmp ,e])
+                     ,body)))
+            (define (inline-field-access k s-rator u-rator args)
+              ;; For imported accessors or for JIT mode, inline the
+              ;; selector with an `unsafe-struct?` test plus `unsafe-struct*-ref`.
+              (define im (hash-ref imports u-rator #f))
+              (define type-id (and (or im for-jitify?)
+                                   (pair? args)
+                                   (null? (cdr args))
+                                   (inline-type-id k im add-import! mutated imports)))
+              (cond
+                [type-id
+                 (define tmp (maybe-tmp (car args) 'v))
+                 (define sel `(if (unsafe-struct? ,tmp ,(schemify type-id))
+                                  (unsafe-struct*-ref ,tmp ,(known-field-accessor-pos k))
+                                  (,s-rator ,tmp)))
+                 (wrap-tmp tmp (car args)
+                           sel)]
+                [else #f]))
+            (define (inline-field-mutate k s-rator u-rator args)
+              (define im (hash-ref imports u-rator #f))
+              (define type-id (and (or im for-jitify?)
+                                   (pair? args)
+                                   (pair? (cdr args))
+                                   (null? (cddr args))
+                                   (inline-type-id k im add-import! mutated imports)))
+              (cond
+                [type-id
+                 (define tmp (maybe-tmp (car args) 'v))
+                 (define tmp-rhs (maybe-tmp (cadr args) 'rhs))
+                 (define mut `(if (unsafe-struct? ,tmp ,(schemify type-id))
+                                  (unsafe-struct*-set! ,tmp ,(known-field-mutator-pos k) ,tmp-rhs)
+                                  (,s-rator ,tmp ,tmp-rhs)))
+                 (wrap-tmp tmp (car args)
+                           (wrap-tmp tmp-rhs (cadr args)
+                                     mut))]
+                [else #f]))
             (or (left-left-lambda-convert rator inline-fuel)
                 (and (positive? inline-fuel)
                      (inline-rator))
                 (let ([s-rator (schemify rator)]
                       [args (map schemify exps)]
                       [u-rator (unwrap rator)])
-                  (let ([plain-app?
+                  (define k (find-known u-rator prim-knowns knowns imports mutated))
-                         (or (known-procedure? (find-known u-rator prim-knowns knowns imports mutated))
+                  (cond
-                             (lambda? rator))])
+                    [(and (not for-cify?)
+                          (known-field-accessor? k)
+                          (inline-field-access k s-rator u-rator args))
+                     => (lambda (e) e)]
+                    [(and (not for-cify?)
+                          (known-field-mutator? k)
+                          (inline-field-mutate k s-rator u-rator args))
+                     => (lambda (e) e)]
+                    [else
+                     (define plain-app? (or (known-procedure? k)
+                                            (lambda? rator)))
                      (left-to-right/app s-rator
                                         args
                                         plain-app? for-cify?
-                                       prim-knowns knowns imports mutated))))]
+                                        prim-knowns knowns imports mutated)])))]
            [`,_
             (let ([u-v (unwrap v)])
               (cond

racket/src/schemify/simple.rkt

@@ -5,7 +5,8 @@
          "import.rkt"
          "mutated-state.rkt")
 
-(provide simple?)
+(provide simple?
+         simple/can-copy?)
 
 ;; Check whether an expression is simple in the sense that its order
 ;; of evaluation isn't detectable. This function receives both
@@ -52,3 +53,18 @@
              (string? e)
              (bytes? e)
              (regexp? e)))])))
+
+(define (simple/can-copy? e prim-knowns knowns imports mutated)
+  (match e
+    [`(quote ,v) (can-copy-literal? v)]
+    [`(,_ . ,_) #f]
+    [`,_
+     (let ([e (unwrap e)])
+       (or (and (symbol? e)
+                (simple-mutated-state? (hash-ref mutated e #f)))
+           (can-copy-literal? e)))]))
+
+(define (can-copy-literal? e)
+  (or (integer? e)
+      (boolean? e)
+      (symbol? e)))