Racket code exploring SSA (Single Static Assignment) form for control flow graphs. Includes a conversion from an arbitrary CFG into SSA, and code to render a CFG as a Racket pict.

ssa.rkt

#lang racket
(require racket/set)
(require graphviz)
(require racket/string)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Simple SSA Conversion and Visualization
;; Michael Vollmer, 2025
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;;; Data structures for our CFG representation

;; Expressions
(struct const (value) #:transparent)
(struct varref (name) #:transparent)
(struct binop (op left right) #:transparent)

;; Statements
(struct assign (var expr) #:transparent)
(struct phi (var args) #:transparent)  ;; args: list of (predecessor-label . var)
(struct if-jump (test then-block else-block) #:transparent)
(struct goto (label) #:transparent)

;; Basic blocks
(struct block (label statements terminator) #:transparent)

;;; CFG representation
(struct cfg (blocks entry) #:transparent)

;;; Create a CFG from a list of blocks
(define (make-cfg blocks entry-label)
  (cfg (make-hash (map (λ (b) (cons (block-label b) b)) blocks)) entry-label))

;;; Graph utilities

(define (get-vertices cfg)
  (hash-keys (cfg-blocks cfg)))

(define (get-successors cfg label)
  (define block (hash-ref (cfg-blocks cfg) label))
  (match (block-terminator block)
    [(goto succ) (list succ)]
    [(if-jump _ then else) (list then else)]
    [_ '()]))

(define (get-predecessors cfg)
  (define preds (make-hash))
  (define vertices (get-vertices cfg))
  
  ;; Initialize with empty sets
  (for ([v vertices]) (hash-set! preds v (set)))
  
  ;; Find predecessors by examining each block's successors
  (for ([v vertices])
    (for ([succ (get-successors cfg v)])
      (hash-set! preds succ (set-add (hash-ref preds succ) v))))
  
  preds)

;;; Dominance analysis implementation

;; Compute dominators using the iterative algorithm
(define (compute-dominators in-cfg)
  (define vertices (get-vertices in-cfg))
  (define entry (cfg-entry in-cfg))
  (define preds (get-predecessors in-cfg))
  (define doms (make-hash))
  
  ;; Initialize
  (for ([v vertices])
    (hash-set! doms v
               (if (equal? v entry)
                   (set v)
                   (list->set vertices))))
  
  ;; Iterate to fixed point
  (let loop ()
    (define changed? #f)
    (for ([v vertices] #:when (not (equal? v entry)))
      (define pred-doms
        (for/list ([p (in-set (hash-ref preds v))])
          (hash-ref doms p)))
      (define new-dom
        (set-union (set v) (set-intersect-all pred-doms)))
      (define old-dom (hash-ref doms v))
      (unless (equal? new-dom old-dom)
        (hash-set! doms v new-dom)
        (set! changed? #t)))
    (when changed? (loop)))
  doms)

;; Compute immediate dominators
(define (compute-immediate-dominators doms)
  (define idoms (make-hash))
  (for ([(node dom-set) (in-hash doms)])
    (let ([candidates (set->list (set-remove dom-set node))]) ; doms except the node itself
      (if (null? candidates)
          (hash-set! idoms node #f)  ;; entry has no immediate dominator
          ;; pick candidate with maximum dominator-set size (deepest one)
          (let ([idom (argmax (λ (d) (set-count (hash-ref doms d))) candidates)])
            (hash-set! idoms node idom)))))
  idoms)


;; Compute dominance frontier
(define (compute-dominance-frontier cfg idoms)
  (define preds (get-predecessors cfg))
  (define df (make-hash))
  (define vertices (get-vertices cfg))

  ;; init DF sets
  (for ([v vertices]) (hash-set! df v (set)))

  ;; for each node b, for each predecessor p of b:
  ;;   walk runner := p up the idom chain until runner == idom[b],
  ;;   adding b to DF(runner) at each step.
  (for ([b vertices])
    (let ([pred-set (hash-ref preds b (set))])
      (when (> (set-count pred-set) 0)
        (let ([idom-b (hash-ref idoms b #f)])
          (for ([p (set->list pred-set)])
            (let loop ([runner p])
              (when (and runner (not (equal? runner idom-b)))
                ;; only update df if runner is a known vertex
                (when (hash-has-key? df runner)
                  (hash-set! df runner (set-add (hash-ref df runner) b)))
                (loop (hash-ref idoms runner #f)))))))))
  df)


;; Helper function for set intersection of multiple sets
(define (set-intersect-all sets)
  (cond
    [(null? sets) (set)]               ; intersection of no sets is empty
    [else (foldl set-intersect (car sets) (cdr sets))]))

;; Simple while loop macro
(define-syntax-rule (while condition body ...)
  (let loop ()
    (when condition
      body ...
      (loop))))

;;; SSA Construction Algorithm

;; Step 1: Find all variable definitions
(define (find-variable-definitions cfg)
  (define blocks (cfg-blocks cfg))
  (define defs (make-hash))
  
  (for ([(label block) (in-hash blocks)])
    (for ([stmt (block-statements block)])
      (when (assign? stmt)
        (let ([var-name (assign-var stmt)])
          (hash-set! defs var-name 
                     (set-add (hash-ref defs var-name (set)) label))))))
  defs)

;; Step 2: Insert Φ functions at dominance frontiers
(define (insert-phi-functions in-cfg)
  (define blocks (cfg-blocks in-cfg))
  (define doms (compute-dominators in-cfg))
  (define idoms (compute-immediate-dominators doms))
  (define df (compute-dominance-frontier in-cfg idoms))
  (define var-defs (find-variable-definitions in-cfg))
  
  ;; For each variable, compute its iterated dominance frontier
  (define phi-placement (make-hash)) ;; block-label -> list of (var . phi-node)
  
  (for ([(var def-blocks) (in-hash var-defs)])
    (when (> (set-count def-blocks) 1)
      (let ([idf (compute-iterated-df df def-blocks)])
        (for ([block-label (in-set idf)])
          (let ([phi-var (string->symbol (format "~a.phi" var))])
            (hash-set! phi-placement block-label
                       (cons (cons var phi-var)
                             (hash-ref phi-placement block-label '()))))))))
  
  ;; Actually insert the Φ functions into the blocks
  (define new-blocks (hash-copy blocks))
  
  (for ([(block-label phi-vars) (in-hash phi-placement)])
    (let* ([old-block (hash-ref blocks block-label)]
           [phi-stmts (for/list ([pv (in-list phi-vars)])
                        (let ([var-name (car pv)]
                              [phi-var (cdr pv)])
                          (phi phi-var '())))]  ;; Empty args for now
           [new-stmts (append phi-stmts (block-statements old-block))])
      (hash-set! new-blocks block-label
                 (block block-label new-stmts (block-terminator old-block)))))
  
  (struct-copy cfg in-cfg [blocks new-blocks]))

;; Compute iterated dominance frontier for a set of blocks
(define (compute-iterated-df df def-blocks)
  (let loop ([worklist (list->set def-blocks)]
             [result (set)])
    (if (set-empty? worklist)
        result
        (let* ([current (set-first worklist)]
               [remaining (set-rest worklist)]
               [current-df (hash-ref df current (set))]
               [new-nodes (set-subtract current-df result)])
          (loop (set-union remaining new-nodes)
                (set-union result new-nodes))))))

;; Step 3: Rename variables with proper Φ argument handling
(define (rename-variables cfg)
  (define blocks (cfg-blocks cfg))
  (define all-vars (find-all-variables cfg))
  (define idoms (compute-immediate-dominators (compute-dominators cfg)))
  
  ;; State for renaming
  (define counters (make-hash)) ;; var -> next version number
  (define stacks (make-hash))   ;; var -> stack of current versions
  (define phi-args (make-hash)) ;; block-label -> list of (predecessor . (var . version))
  
  ;; Initialize counters and stacks
  (for ([var all-vars])
    (hash-set! counters var 0)
    (hash-set! stacks var '()))
  
  ;; Recursive function to process blocks in dominator tree order
  (define (process-block label)
    (define new-block (hash-ref blocks label))
    (define new-stmts '())
    (define renamed-stmts '())
    
    ;; First, process Φ functions (they're definitions)
    (for ([stmt (block-statements new-block)])
      (if (phi? stmt)
          (let* ([phi-var (phi-var stmt)]
                 [base-var (get-base-var phi-var)]
                 [new-version (new-version! base-var counters)])
            ;; Push the new version and rename the Φ variable
            (push-version! stacks base-var new-version)
            (set! new-stmts (cons (phi new-version '()) new-stmts)))
          ;; accumulate the non-phi stmt onto renamed-stmts
          (set! renamed-stmts (cons stmt renamed-stmts))))
    
    ;; Now process regular statements in reverse order (we'll reverse later)
    (set! renamed-stmts (reverse renamed-stmts))
    (let loop ([stmts renamed-stmts] [acc '()])
      (if (null? stmts)
          (set! new-stmts (append (reverse acc) new-stmts))
          (let* ([stmt (car stmts)]
                 [rest (cdr stmts)])
            (match stmt
              [(assign var expr)
               (let* ([base-var (get-base-var var)]
                      [new-var (new-version! base-var counters)]
                      [new-expr (rename-expr expr stacks)])
                 (push-version! stacks base-var new-var)
                 (loop rest (cons (assign new-var new-expr) acc)))]
              [else
               (loop rest (cons (rename-stmt stmt stacks) acc))]))))
    
    ;; Update the block with renamed statements
    (hash-set! blocks label 
               (block label new-stmts (block-terminator new-block)))
    
    ;; Process children in dominator tree (blocks that this block immediately dominates)
    (for ([child (get-vertices cfg)])
      (when (equal? (hash-ref idoms child) label)
        (process-block child)))
    
    ;; Before leaving the block, record Φ arguments for successors
    (record-phi-arguments cfg stacks phi-args label)
    
    ;; Pop versions we pushed in this block
    (pop-block-versions stacks (block-statements new-block)))
  
  ;; Start processing from entry block
  (process-block (cfg-entry cfg))
  
  ;; Now update the Φ functions with their arguments
  (update-phi-functions cfg phi-args))

;; Helper functions for renaming
(define (new-version! var counters)
  (let ([count (hash-ref counters var 0)])
    (hash-set! counters var (+ count 1))
    (string->symbol (format "~a.~a" var count))))

(define (push-version! stacks var version)
  (hash-set! stacks var (cons version (hash-ref stacks var '()))))

(define (pop-block-versions stacks stmts)
  (for ([stmt stmts])
    (when (or (assign? stmt) (phi? stmt))
      (let* ([var (if (assign? stmt) (assign-var stmt) (phi-var stmt))]
             [base-var (get-base-var var)])
        (when (hash-has-key? stacks base-var)
          (let ([stack (hash-ref stacks base-var)])
            (when (not (null? stack))
              (hash-set! stacks base-var (cdr stack)))))))))

(define (get-base-var var)
  (let ([str (symbol->string var)])
    (if (string-contains? str ".")
        (string->symbol (car (string-split str ".")))
        var)))

(define (rename-expr expr stacks)
  (match expr
    [(const value) expr]
    [(binop op left right)
     (binop op (rename-expr left stacks) (rename-expr right stacks))]
    [(varref name)
     (let* ([base-var (get-base-var name)]
            [stack (hash-ref stacks base-var '())])
       (if (null? stack)
           (varref name)  ;; Use original name if no version on stack
           (varref (car stack))))]
    [else expr]))

(define (rename-stmt stmt stacks)
  (match stmt
    [(if-jump test then else)
     (if-jump (rename-expr test stacks) then else)]
    [else stmt]))

;; Record Φ arguments for successor blocks
(define (record-phi-arg-for-successor cfg phi-args pred-label succ-label stacks)
  (define succ-block (hash-ref (cfg-blocks cfg) succ-label))
  (for ([stmt (block-statements succ-block)])
    (when (phi? stmt)
      (let* ([phi-var (phi-var stmt)]
             [base-var (get-base-var phi-var)]
             [stack (hash-ref stacks base-var '())])
        (when (not (null? stack))
          (let ([current-version (car stack)])
            ;; store a proper list [pred base-var version] for easy filtering later
            (hash-set! phi-args succ-label
                       (cons (list pred-label base-var current-version)
                             (hash-ref phi-args succ-label '())))))))))

(define (record-phi-arguments cfg stacks phi-args current-label)
  (define current-block (hash-ref (cfg-blocks cfg) current-label))
  (define terminator (block-terminator current-block))
  
  (match terminator
    [(goto succ-label)
     (record-phi-arg-for-successor cfg phi-args current-label succ-label stacks)]
    [(if-jump test then-label else-label)
     (record-phi-arg-for-successor cfg phi-args current-label then-label stacks)
     (record-phi-arg-for-successor cfg phi-args current-label else-label stacks)]
    [_ (void)]))

;; Update Φ functions with their arguments
(define (update-phi-functions in-cfg phi-args)
  (define blocks (cfg-blocks in-cfg))
  (define new-blocks (make-hash))
  
  (for ([(label this-block) (in-hash blocks)])
    (let* ([new-stmts 
            (for/list ([stmt (block-statements this-block)])
              (if (phi? stmt)
                  (let* ([phi-var (phi-var stmt)]
                         [base-var (get-base-var phi-var)]
                         [args (hash-ref phi-args label '())]
                         ;; select only entries for this base-var
                         [var-args (filter (λ (triple) (equal? (second triple) base-var)) args)]
                         ;; format as (pred . version)
                         [formatted-args (map (λ (triple)
                                                (cons (first triple) (third triple)))
                                              var-args)])
                    (phi phi-var formatted-args))
                  stmt))])
      (hash-set! new-blocks label 
                 (block label new-stmts (block-terminator this-block)))))
  
  (struct-copy cfg in-cfg [blocks new-blocks]))

;; Find all variables in the CFG
(define (find-all-variables cfg)
  (define blocks (cfg-blocks cfg))
  (define vars (mutable-set))
  
  (for ([(label block) (in-hash blocks)])
    (for ([stmt (block-statements block)])
      (when (assign? stmt)
        (set-add! vars (assign-var stmt)))
      (extract-vars-from-expr 
       (if (assign? stmt) (assign-expr stmt) 
           (if (if-jump? stmt) (if-jump-test stmt) #f)) 
       vars)))
  (set->list vars))

(define (extract-vars-from-expr expr vars)
  (when expr
    (match expr
      [(var name) (set-add! vars name)]
      [(binop op left right)
       (extract-vars-from-expr left vars)
       (extract-vars-from-expr right vars)]
      [_ (void)])))

;;; Main SSA transformation function
(define (cfg-to-ssa in-cfg)
  (rename-variables (insert-phi-functions in-cfg)))

; Small pretty-printer for expressions / statements used in labels.
(define (expr->str e)
  (cond
    [(void? e) ""]
    [(const? e) (format "~a" (const-value e))]
    [(varref? e) (symbol->string (varref-name e))]
    [(binop? e)
     (format "(~a ~a ~a)"
             (binop-op e)
             (expr->str (binop-left e))
             (expr->str (binop-right e)))]
    [else (format "~a" e)]))

(define (phi-args->str args)
  ;; args: list of (pred . var) pairs
  (string-join
   (map (λ (pair) (format "~a:~a" (car pair) (expr->str (cdr pair))))
        args)
   ", "))

(define (stmt->str s)
  (cond
    [(assign? s) (format "assign(~a, ~a)" (symbol->string (assign-var s)) (expr->str (assign-expr s)))]
    [(phi? s)    (format "assign(~a, Φ(~a))" (symbol->string (phi-var s)) (phi-args->str (phi-args s)))]
    [(if-jump? s) (format "if ~s then ~s else ~s"
                         (expr->str (if-jump-test s)) (if-jump-then-block s) (if-jump-else-block s))]
    [(goto? s) (format "goto ~s" (goto-label s))]
    [else (format "~a" s)]))

;; Build a human-friendly label: φs first, then ordinary statements, then the terminator
(define (block->label-string block)
  (define stmts (block-statements block))
  (define phis (filter phi? stmts))
  (define others (filter (λ (x) (not (phi? x))) stmts))
  (define phi-lines (map stmt->str phis))
  (define stmt-lines (map stmt->str others))
  (define term-line
    (match (block-terminator block)
      [(goto lbl) (format "→ ~a" lbl)]
      [(if-jump test then else) (format "→ if ~a then ~a else ~a" (expr->str test) then else)]
      [else (format "→ ~a" (block-terminator block))]))
  (string-join (append phi-lines stmt-lines (list term-line)) "\n"))

;; Render a CFG as a Graphviz pict
;; Main renderer
(define (render-cfg->pict cfg)
  (define blocks (cfg-blocks cfg))
  (define entry (cfg-entry cfg))

  ;; nodes: a list of (label #:label <escaped-string> #:shape "box" ...)
  (define nodes
    (for/list ([(lbl b) (in-hash blocks)])
      `(,lbl
        #:label ,(block->label-string b)
        #:xlabel ,lbl
        #:shape "box"
        #:style "filled"
        #:fillcolor ,(if (equal? lbl entry) "lightgray" "gray"))))

  ;; edges as "a -> b" strings
  (define edges
    (for*/list ([(lbl b) (in-hash blocks)]
                [succ (in-list (get-successors cfg lbl))])
      (format "~a -> ~a" lbl succ)))
  (define graph-structure (append nodes edges))
  (digraph->pict (make-digraph graph-structure)))

;;; Example usage with proper Φ nodes

(define example-blocks
  (list
   (block "entry"
          (list (assign 'x (const 1)))
          (if-jump (varref 'condition) "then" "else"))
   
   (block "then" 
          (list (assign 'x (binop '+ (varref 'x) (const 1))))
          (goto "exit"))
   
   (block "else"
          (list (assign 'x (const 2)))
          (goto "exit"))
   
   (block "exit"
          (list (assign 'y (varref 'x)))
          (goto "end"))
   
   (block "end" '() (goto "end"))))

(define loop-example-blocks
  (list
   ;; Entry: initialize i = 0
   (block "entry"
          (list (assign 'i (const 0)))
          (goto "loop"))

   ;; Loop header: φ for i
   (block "loop"
          '() ;; φ for i will get inserted here
          (if-jump (binop '< (varref 'i) (const 10))
                   "body"
                   "exit"))

   ;; Loop body: increment i
   (block "body"
          (list (assign 'i (binop '+ (varref 'i) (const 1))))
          (goto "loop"))

   ;; Exit block
   (block "exit"
          (list (assign 'y (varref 'i)))
          (goto "end"))

   ;; Dummy end block
   (block "end" '() (goto "end"))))


(define example-cfg (make-cfg loop-example-blocks "entry"))

(define (test-proper-ssa)
  (displayln "Original CFG:")
  (for ([(label block) (in-hash (cfg-blocks example-cfg))])
    (printf "Block ~a:\n" label)
    (for ([stmt (block-statements block)])
      (printf "  ~a\n" stmt))
    (printf "  -> ~a\n\n" (block-terminator block)))
  
  (displayln "Computing dominators...")
  (define doms (compute-dominators example-cfg))
  (for ([(v dom-set) (in-hash doms)])
    (printf "DOM(~a) = ~a\n" v (set->list dom-set)))
  
  (displayln "\nComputing immediate dominators...")
  (define idoms (compute-immediate-dominators doms))
  (for ([(v idom) (in-hash idoms)])
    (printf "IDOM(~a) = ~a\n" v idom))
  
  (displayln "\nComputing dominance frontier...")
  (define df (compute-dominance-frontier example-cfg doms))
  (for ([(v df-set) (in-hash df)])
    (printf "DF(~a) = ~a\n" v (set->list df-set)))
  
  (displayln "\nAfter SSA transformation with proper Φ nodes:")
  (define ssa-cfg (cfg-to-ssa example-cfg))
  (for ([(label block) (in-hash (cfg-blocks ssa-cfg))])
    (printf "Block ~a:\n" label)
    (for ([stmt (block-statements block)])
      (match stmt
        [(phi var args)
         (printf "  ~a = Φ(" var)
         (for ([arg args] [i (in-naturals)])
           (when (> i 0) (printf ", "))
           (printf "~a: ~a" (car arg) (cdr arg)))
         (printf ")\n")]
        [else (printf "  ~a\n" stmt)]))
    (printf "  -> ~a\n\n" (block-terminator block))))

;; Run the test
(test-proper-ssa)

(displayln "Original CFG:")
(render-cfg->pict example-cfg)
(displayln "CFG in SSA form:")
(render-cfg->pict (cfg-to-ssa example-cfg))