Skip to content

Instantly share code, notes, and snippets.

@dannypsnl
Last active December 17, 2025 08:25
Show Gist options
  • Select an option

  • Save dannypsnl/9be2f44fcb29d0ce1116123cda9f13b4 to your computer and use it in GitHub Desktop.

Select an option

Save dannypsnl/9be2f44fcb29d0ce1116123cda9f13b4 to your computer and use it in GitHub Desktop.
egg first attempt and visualization
#lang racket
(require (for-syntax syntax/parse))
(require data/queue)
(require graphviz)
(struct e-class (e-nodes))
(struct e-graph (e-classes) #:transparent)
#|
The rewriter is a function that takes
1. a callback (to add result into same e-class), and
2. a e-node N
if N is matched by the pattern then we add rewrite result to e-class
|#
(define-syntax (=> stx)
(syntax-parse stx
[(_ pat rewrite-result)
#'(lambda (add! e-node)
(match e-node
[pat (add! rewrite-result)]
[else (void)]))]))
#| Language to operate |#
(struct Add (a b) #:transparent)
(struct Mul (a b) #:transparent)
(struct Div (a b) #:transparent)
(struct ShiftL (a b) #:transparent)
#| Rewrites of the language |#
; FIXME: this way won't work, because the e-node 1 is in a e-class
; and we have some candidate e-nodes by sequence to work on,
; so this naive match cannot handle the idea
(define R1 (=> (Mul x 1) x))
(define R2 (=> (Div (Mul x y) z) (Mul x (Div y z))))
#| helpers |#
(define (num g k)
(define c (e-class (make-queue)))
(enqueue! (e-class-e-nodes c) k)
(enqueue! (e-graph-e-classes g) c)
c)
(define (var g x)
(define c (e-class (make-queue)))
(enqueue! (e-class-e-nodes c) x)
(enqueue! (e-graph-e-classes g) c)
c)
(define (mul g a b)
(define c (e-class (make-queue)))
(enqueue! (e-class-e-nodes c) (Mul a b))
(enqueue! (e-graph-e-classes g) c)
c)
(define (div g a b)
(define c (e-class (make-queue)))
(enqueue! (e-class-e-nodes c) (Div a b))
(enqueue! (e-graph-e-classes g) c)
c)
(define g (e-graph (make-queue)))
(define two (num g 2))
(define target (div g (mul g (var g 'a) two) two))
#;(eq? target two) ; #f by reference to different e-class
#;(eq? two two) ; #t by reference to same
#| Visualize |#
(define name-map (make-hasheq))
(define (node-name en)
(define id (hash-ref name-map en #f))
(if id id
(let ([id (symbol->string (gensym))])
(hash-set! name-map en id)
id)))
(define (build-node en)
(define id (node-name en))
(match en
[(? number? k) (list id '#:label (number->string k))]
[(? symbol? x) (list id '#:label (symbol->string x))]
[(Add a b) (list id '#:shape "diamond" '#:label "+")]
[(Mul a b) (list id '#:shape "diamond" '#:label "*")]
[(Div a b) (list id '#:shape "diamond" '#:label "/")]
[(ShiftL a b) (list id '#:shape "diamond" '#:label "<<")]))
(define (node-arrow arrows en)
(match en
[(Add a b)
(enqueue! arrows (string-append (node-name en) " -> " (node-name a)))
(enqueue! arrows (string-append (node-name en) " -> " (node-name b)))]
[(Mul a b)
(enqueue! arrows (string-append (node-name en) " -> " (node-name a)))
(enqueue! arrows (string-append (node-name en) " -> " (node-name b)))]
[(Div a b)
(enqueue! arrows (string-append (node-name en) " -> " (node-name a)))
(enqueue! arrows (string-append (node-name en) " -> " (node-name b)))]
[(ShiftL a b)
(enqueue! arrows (string-append (node-name en) " -> " (node-name a)))
(enqueue! arrows (string-append (node-name en) " -> " (node-name b)))]
[else (void)]))
(define (visualize g)
(define classes
(for/list ([ec (queue->list (e-graph-e-classes g))])
(define id (node-name ec))
(define nodes
(for/list ([en (queue->list (e-class-e-nodes ec))]) (build-node en)))
(list 'subgraph id nodes)))
(define arrows (make-queue))
(for* ([ec (queue->list (e-graph-e-classes g))]
[en (queue->list (e-class-e-nodes ec))])
(node-arrow arrows en))
(make-digraph
`(,@classes
,@(queue->list arrows))))
#| Start |#
(define (grow egraph rules)
(for ([_ (in-range 1000)])
(for ([R rules])
(for ([ec (queue->list (e-graph-e-classes egraph))])
(define q (e-class-e-nodes ec))
(for ([enode (queue->list q)])
(R (lambda (newnode)
(enqueue! q newnode))
enode)))))
; TODO: how to extract the best?
egraph)
(digraph->pict (visualize g))
(define g- (grow g (list R1 R2)))
(digraph->pict (visualize g-))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment