## Representation and JITting

Let's take a closer look at the `sphere` example from the Ao project page:

``````(define (sphere r)
(lambda ( x y z )
(- (sqrt (+ (* x x) (* y y) (* z z))) r)))
``````

Looking past the blinding walls of parenthesis, this is implementing the function
`f(x, y, z) = sqrt(x**2 + y**2 + z**2) - r`

Ao implements an f-rep kernel, where shapes are functions of `(x, y, z)` that return a single value: negative inside the model, positive outside.

This representation dovetails nicely with Scheme as a user-friendly scripting language. For example, the union of two shapes is the minimum of their two functions evaluated at the target point:

``````(define (union a b)
(lambda (x y z)
(min (a x y z) (b x y z))))
``````

Lambda functions are not well-optimized for evaluation in an f-rep kernel, so they're transformed into a more efficient internal representation.

This transformation uses a strategy similar to that of a tracing JIT.

The function is called with special objects that record the sequence of arithmetic operations that they experience, rebuilding the arithmetic tree by tracing the computation.

Let's walk through a basic implementation in pure Scheme.

``````;; A token is a list with an operation and
;; some number of arguments
(define (make-token opcode . args)
(cons opcode args))

;; Test to see if the given object is a token
(define (token? t)
(and (list? t) (symbol? (car t))))

;; wrap-token is a no-op on tokens and
;; wraps numbers in a constant token
(define (wrap-token t)
(cond ((token? t) t)
((number? t) (make-token 'constant t))
(else (error "Can't wrap" t))))
``````

Here's the most basic implementation of the compiler:

``````(define (jit f)
(f (make-token 'x) (make-token 'y) (make-token 'z)))
``````

This evaluation needs to happen in an environment which understands how to combine tokens, so all of the relevant operations are overloaded. Here's an implementation of overloading for addition:

``````(define +_ +)   ; save the existing addition operator
(define (+ a b) ; then overload it with a token-aware replacement
(if (or (token? a) (token? b))
(make-token 'add (wrap-token a) (wrap-token b))
(+_ a b)))
``````

Testing out this implementation, it correctly converts a lambda function back into its original tree of operations:

``````scheme@(guile-user)> (define (f x y z) (+ (+ x y) 2))
scheme@(guile-user)> (jit f)
\$9 = (add (add (x) (y)) (constant 2))
``````

At first, this seems like a silly thing to do: if we already have all of the original s-expressions, why not convert them directly into the tree? The tracing strategy gives you the rest of Scheme's syntax for free.

Consider the following function:

``````(define (circle xmin ymin d)
" Define a circle from its lower-left corner and diameter "
(let* ((r (/ d 2))
(x0 (+ xmin r))
(y0 (+ ymin r))
(square (lambda (a) (* a a))))
(lambda (x y z)
(- (sqrt (+ (square (- x x0))
(square (- y y0)))) r))))
``````

When evaluated with this tracing JIT, Scheme handles all of the arithmetic and variable bindings in the `let` expression; doing it ourselves would involve re-implementing all of that ourselves (with something like the meta-circular evaluator).

The actual JIT is more complicated than the sample code above. In particular, it properly handles n-ary operations (i.e. `(+ x y z 2)`) and drops identity operations (i.e. `(+ x 0)`).

The JIT also applies a handful of optimizations when generating the tree: