Add expanded identifier table utility (#710)

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: jackfirth <[email protected]>

Author: Copilot

default-recommendations/analyzers/identifier-usage.rkt

@@ -9,13 +9,11 @@
   [identifier-usage-analyzer expansion-analyzer?]))
 
 
-(require racket/hash
-         racket/list
-         racket/set
+(require racket/list
          racket/stream
-         racket/treelist
-         rebellion/collection/hash
+         rebellion/collection/entry
          rebellion/streaming/transducer
+         resyntax/default-recommendations/analyzers/private/expanded-id-table
          resyntax/private/analyzer
          resyntax/private/syntax-path
          resyntax/private/syntax-property-bundle
@@ -34,17 +32,6 @@
 ;@----------------------------------------------------------------------------------------------------
 
 
-(define (append-all-id-maps id-maps)
-  (for/fold ([combined (hash)])
-            ([map id-maps])
-    (hash-union combined map #:combine treelist-append)))
-
-
-(define (id-map-shift-phase id-map levels)
-  (for/hash ([(phase ids) (in-hash id-map)])
-    (values (and phase (+ phase levels)) ids)))
-
-
 ;; Label syntax with phase information
 (define (syntax-label-id-phases expanded-stx)
   (let loop ([expanded-stx expanded-stx] [phase 0] [skip? #false])
@@ -166,56 +153,41 @@
        (stream (attribute id))])))
 
 
-(define (phase-binding-table bound-ids used-ids #:phase phase)
-  (define initial-map
-    (for/fold ([map (make-immutable-free-id-table #:phase phase)])
-              ([bound bound-ids])
-      (free-id-table-set map bound '())))
-  (for*/fold ([map initial-map])
-             ([bound bound-ids]
-              [used used-ids]
-              #:when (free-identifier=? bound used))
-    (free-id-table-update map bound (λ (previous) (cons used previous)) '())))
-
-
-(define (identifier-binding-table bound-ids-by-phase used-ids-by-phase)
-  (for/hash
-      ([phase
-        (in-set (set-union (hash-key-set bound-ids-by-phase) (hash-key-set used-ids-by-phase)))])
-    (define bound-ids (hash-ref bound-ids-by-phase phase '()))
-    (define used-ids (hash-ref used-ids-by-phase phase '()))
-    (values phase (phase-binding-table bound-ids used-ids #:phase phase))))
-
-
 (define (fully-expanded-syntax-binding-table stx)
   (define labeled-stx (syntax-label-id-phases (syntax-label-paths stx 'expanded-path)))
 
-  ;; Get bound identifiers and group by phase
-  (define bound-ids-by-phase
-    (for/fold ([result (hash)])
-              ([id (in-stream (binding-site-identifiers labeled-stx))])
-      (define id-phase (syntax-property id 'phase))
-      (hash-update result id-phase (λ (prev) (treelist-add prev id)) (treelist))))
+  ;; Create expanded-id-table to track bound identifiers with empty usage lists
+  (define table (make-expanded-id-table))
+  
+  ;; Initialize all bound identifiers with empty usage lists
+  (for ([id (in-stream (binding-site-identifiers labeled-stx))])
+    (define id-phase (syntax-property id 'phase))
+    (expanded-id-table-set! table (expanded-identifier id id-phase) '()))
 
-  ;; Get used identifiers and group by phase
-  (define used-ids-by-phase
-    (for/fold ([result (hash)])
-              ([id (in-stream (usage-site-identifiers labeled-stx))])
-      (define id-phase (syntax-property id 'phase))
-      (hash-update result id-phase (λ (prev) (treelist-add prev id)) (treelist))))
+  ;; For each usage, find its binding within the same phase and add it to the usage list
+  (for ([used-id (in-stream (usage-site-identifiers labeled-stx))])
+    (define used-phase (syntax-property used-id 'phase))
+    (for ([bound-entry (in-expanded-id-table-phase table used-phase)])
+      (define bound-expanded-id (entry-key bound-entry))
+      (define bound-id (expanded-identifier-syntax bound-expanded-id))
+      (when (free-identifier=? bound-id used-id)
+        (define current-usages (entry-value bound-entry))
+        (expanded-id-table-set! table bound-expanded-id (cons used-id current-usages)))))
 
-  (identifier-binding-table bound-ids-by-phase used-ids-by-phase))
+  table)
 
 
 (define identifier-usage-analyzer
   (make-expansion-analyzer
    #:name 'identifier-usage-analyzer
    (λ (expanded-stx)
      (define table (fully-expanded-syntax-binding-table expanded-stx))
-     (transduce (in-hash-values table)
-                (append-mapping
-                 (λ (id-table)
-                   (for/stream ([(bound-id usages) (in-free-id-table id-table)])
-                     (define exp-path (syntax-property bound-id 'expanded-path))
-                     (syntax-property-entry exp-path 'usage-count (length usages)))))
+     (transduce (in-expanded-id-table table)
+                (mapping
+                 (λ (entry)
+                   (define expanded-id (entry-key entry))
+                   (define usages (entry-value entry))
+                   (define bound-id (expanded-identifier-syntax expanded-id))
+                   (define exp-path (syntax-property bound-id 'expanded-path))
+                   (syntax-property-entry exp-path 'usage-count (length usages))))
                 #:into into-syntax-property-bundle))))

default-recommendations/analyzers/private/expanded-id-table.rkt

@@ -0,0 +1,121 @@
+#lang racket/base
+
+
+(require racket/contract/base)
+
+
+(provide
+ (struct-out expanded-identifier)
+ (contract-out
+  [expanded-id-table? (-> any/c boolean?)]
+  [make-expanded-id-table (-> expanded-id-table?)]
+  [expanded-id-table-ref
+   (->* (expanded-id-table? expanded-identifier?) (failure-result/c) any/c)]
+  [expanded-id-table-set! (-> expanded-id-table? expanded-identifier? any/c void?)]
+  [in-expanded-id-table
+   (-> expanded-id-table? (sequence/c (entry/c expanded-identifier? any/c)))]
+  [in-expanded-id-table-phase
+   (-> expanded-id-table? (or/c exact-nonnegative-integer? #false) (sequence/c (entry/c expanded-identifier? any/c)))]))
+
+
+(require guard
+         racket/contract/base
+         racket/dict
+         racket/match
+         racket/sequence
+         racket/stream
+         rebellion/base/result
+         rebellion/collection/entry
+         syntax/id-table)
+
+
+(module+ test
+  (require (submod "..")
+           rackunit))
+
+
+;@----------------------------------------------------------------------------------------------------
+
+
+(struct expanded-identifier (syntax phase)
+  #:transparent
+  #:guard (struct-guard/c identifier? (or/c exact-nonnegative-integer? #false)))
+
+
+(struct expanded-id-table (table))
+
+
+(define (make-expanded-id-table)
+  (expanded-id-table (make-hasheq)))
+
+
+(define (expanded-id-table-ref table id [failure-result (λ () (error 'expanded-id-table-ref "no mapping for ~a" id))])
+  (define phase (expanded-identifier-phase id))
+  (define stx (expanded-identifier-syntax id))
+  (define phase-table (hash-ref (expanded-id-table-table table) phase #false))
+  (if phase-table
+      (free-id-table-ref phase-table stx failure-result)
+      (if (procedure? failure-result)
+          (failure-result)
+          failure-result)))
+
+
+(define (expanded-id-table-set! table id value)
+  (define phase (expanded-identifier-phase id))
+  (define stx (expanded-identifier-syntax id))
+  (define phase-table
+    (hash-ref! (expanded-id-table-table table)
+               phase
+               (λ () (make-free-id-table #:phase phase))))
+  (free-id-table-set! phase-table stx value))
+
+
+(define (in-expanded-id-table table)
+  (for*/stream ([(phase phase-table) (in-hash (expanded-id-table-table table))]
+                [(stx value) (in-free-id-table phase-table)])
+    (entry (expanded-identifier stx phase) value)))
+
+
+(define (in-expanded-id-table-phase table phase)
+  (define phase-table (hash-ref (expanded-id-table-table table) phase #false))
+  (if phase-table
+      (for/stream ([(stx value) (in-free-id-table phase-table)])
+        (entry (expanded-identifier stx phase) value))
+      (stream)))
+
+
+;@----------------------------------------------------------------------------------------------------
+
+
+(module+ test
+  (test-case "expanded-id-table"
+    
+    (test-case "make-expanded-id-table creates empty table"
+      (define table (make-expanded-id-table))
+      (check-pred expanded-id-table? table))
+    
+    (test-case "expanded-id-table-set! and expanded-id-table-ref"
+      (define table (make-expanded-id-table))
+      (define id1 (expanded-identifier #'x 0))
+      (expanded-id-table-set! table id1 'foo)
+      (check-equal? (expanded-id-table-ref table id1) 'foo))
+    
+    (test-case "identifiers at different phases are distinct"
+      (define table (make-expanded-id-table))
+      (define id-phase0 (expanded-identifier #'x 0))
+      (define id-phase1 (expanded-identifier #'x 1))
+      (expanded-id-table-set! table id-phase0 'phase0-value)
+      (expanded-id-table-set! table id-phase1 'phase1-value)
+      (check-equal? (expanded-id-table-ref table id-phase0) 'phase0-value)
+      (check-equal? (expanded-id-table-ref table id-phase1) 'phase1-value))
+    
+    (test-case "in-expanded-id-table iterates over entries"
+      (define table (make-expanded-id-table))
+      (define id1 (expanded-identifier #'x 0))
+      (define id2 (expanded-identifier #'y 0))
+      (define id3 (expanded-identifier #'z 1))
+      (expanded-id-table-set! table id1 'val1)
+      (expanded-id-table-set! table id2 'val2)
+      (expanded-id-table-set! table id3 'val3)
+      (define entries (for/list ([e (in-expanded-id-table table)]) e))
+      (check-equal? (length entries) 3))))

default-recommendations/analyzers/variable-mutability.rkt

@@ -9,11 +9,12 @@
   [variable-mutability-analyzer expansion-analyzer?]))
 
 
-(require racket/dict
-         racket/list
+(require racket/list
          racket/match
          racket/stream
+         rebellion/collection/entry
          rebellion/streaming/transducer
+         resyntax/default-recommendations/analyzers/private/expanded-id-table
          resyntax/private/analyzer
          resyntax/private/syntax-path
          resyntax/private/syntax-property-bundle
@@ -104,20 +105,19 @@
 
 (define (variable-mutability stx)
   (define labeled-stx (syntax-label-id-phases (syntax-label-paths stx 'expanded-path)))
-  (define variable-table (make-hash))
+  (define variable-table (make-expanded-id-table))
   (for ([id (in-stream (binding-site-variables labeled-stx))])
     (define phase (syntax-property id 'phase))
-    (define phase-specific-table
-      (hash-ref! variable-table phase (λ () (make-free-id-table #:phase phase))))
-    (free-id-table-set! phase-specific-table id 'immutable))
+    (expanded-id-table-set! variable-table (expanded-identifier id phase) 'immutable))
   (for ([id (in-stream (mutated-variables labeled-stx))])
-    (define phase-specific-table (hash-ref variable-table (syntax-property id 'phase)))
-    (free-id-table-set! phase-specific-table id 'mutable))
-  (transduce (in-hash-values variable-table)
-             (append-mapping in-dict-pairs)
+    (define phase (syntax-property id 'phase))
+    (expanded-id-table-set! variable-table (expanded-identifier id phase) 'mutable))
+  (transduce (in-expanded-id-table variable-table)
              (mapping
               (λ (e)
-                (match-define (cons id mode) e)
+                (define expanded-id (entry-key e))
+                (define mode (entry-value e))
+                (define id (expanded-identifier-syntax expanded-id))
                 (define path (syntax-property id 'expanded-path))
                 (syntax-property-entry path 'variable-mutability mode)))
              #:into into-syntax-property-bundle))