An attempt at faking dynamicity by wrapping Coalton

coalton-jit-helper.lisp

(in-package #:coalton-impl/entry)

(defun direct-application-entry-point (node)
  (declare (type parser:node node))

  (let ((env *global-environment*))

    (multiple-value-bind (ty preds accessors node subs)
        (tc:infer-expression-type (tc:resolve-control-flow
                                   (parser:rename-variables node))
                                  (tc:make-variable :kind tc:+kstar+ :allow-result-p t)
                                  nil
                                  (tc:make-tc-env :env env))

      (multiple-value-bind (preds subs)
          (tc:solve-fundeps env preds subs)

        (setf accessors (tc:apply-substitution subs accessors))

        (multiple-value-bind (accessors subs_)
            (tc:solve-accessors accessors env)
          (setf subs (tc:compose-substitution-lists subs subs_))

          (when accessors
            (tc:tc-error "Ambiguous accessor"
                         (source:note (first accessors)
                                      "accessor is ambiguous")))

          (let* ((preds (tc:reduce-context env preds subs))
                 (subs (tc:compose-substitution-lists
                        (tc:default-subs env nil preds)
                        subs))
                 (preds (tc:reduce-context env preds subs))

                 (node (tc:apply-substitution subs node))
                 (ty (tc:apply-substitution subs ty))

                 (qual-ty (tc:qualify preds ty))
                 (scheme (tc:quantify (tc:type-variables qual-ty) qual-ty)))

            (when (null preds)
              (return-from direct-application-entry-point
                (let ((node (codegen:optimize-node
                             (codegen:translate-expression node nil env)
                             env)))
                  (codegen:direct-application
                   node
                   (codegen:make-function-table env)))))

            (tc:tc-error "Unable to codegen"
                         (tc:tc-note node
                                     "expression has ambiguous type ~A"
                                     (tc:type-to-string scheme env))
                         (tc:tc-note node
                                     "Add a type assertion with THE to resolve ambiguity"))))))))

coalton-jit.org

Just-In-Time Compilation for Coalton

I’m looking for a certain language. Common Lisp is one. But it’s more of a tool for constructing other languages. Coalton is one attempt. But it is statically typed. But amongst other things, I want dynamicity. Here’s an attempt at how one can fake dynamicity by wrapping Coalton.

What do I want to do? I want (to write) a number crunching and numerical computing library. This leads me ask for several features:

1. Polymorphic: By polymorphic, I mean the same code can be used for different data types - (signed-byte 16) instead of fixnum, or single-float instead of double-float. I don’t want to be tied down by others’ choices of data types. Similarly, I don’t want to tie up someone by my choices of data types. Thus, the code should be agnostic about concrete data types, that is, it should be polymorphic, at least to some extent. With Common Lisp, there are two main ways to achieve polymorphicity: generic-functions and typecase-variants.
2. Speed-Optimizable: Generic functions are terrible for fast number crunching by themselves.
3. Extensible: Typecase-variants are closed, in that, once the code is written it is not extensible to new types. (Though, we also have Peltadot.)
4. Dynamic: It turns out that Coalton has the above properties. But it is statically typed. To call most of the coalton functions from lisp, and especially polymorphic ones, one must wrap them inside (coalton ...). For coalton users, this is not a problem. They can even use a special REPL and evaluator/compiler that does it automatically for them. But must we expect every lisp user to also be familiar with Coalton? I don’t think that is a good idea. Lisp already has plenty to overwhelm a new programmer without even bringing in Emacs.

Here’s an example of what I mean by being required to wrap inside (coalton ...):

(cl:defpackage :coal-lib
  (:use :coalton
        :coalton-prelude)
  (:export #:add))

(cl:in-package :coal-lib)

(coalton-toplevel
  (define (add x y)
    (+ x y)))

(cl:defpackage :coal-lib-user
  (:use :cl))

(cl:in-package :coal-lib-user)

;; errors because coal-lib:add as a lisp function takes three arguments
(coal-lib:add 2 3)

;; works if you wrap
(coalton:coalton (coal-lib:add 2 3)) ;=> 5
(coalton:coalton (coal-lib:add 2 3.0)) ;=> 5.0 

One option to side-step wrapping is by monomorphize-ing the function. For example:

(coalton-toplevel
  (declare add/ifix/ifix (ifix * ifix -> ifix))
  (define (add/ifix/ifix x y)
    (+ x y)))

But, we can do this automatically. Let us also define a coal-lib-exported package that exports the wrapper function.

(cl:defpackage :coal-lib-exported
  (:use :cl)
  (:export #:add))

(in-package :coal-lib-exported)

(defgeneric coalton-type-of (object))

(defmethod coalton-type-of ((object float))
  (etypecase object
    (single-float 'coalton:single-float)
    (double-float 'coalton:double-float)))

(defmethod coalton-type-of ((object integer))
  (etypecase object
    (fixnum 'coalton:ifix)
    (integer 'coalton:integer)))

(let ((cache (make-hash-table :test #'equal)))
  (defun add (x y)
    (let ((types (list (coalton-type-of x) (coalton-type-of y))))
      (funcall (alexandria:ensure-gethash
                types
                cache
                (let ((name (apply #'alexandria:symbolicate 'add
                                   (loop :for type :in types
                                         :nconcing (list '/ type))))
                      (input-type (cons (car types)
                                        (loop :for type :in (cdr types)
                                              :nconcing (list '* type))))
                      ;; assume output type is the same as first of input type
                      (output-type (first types)))
                  (eval (print 
                         `(coalton:coalton-toplevel
                            (coalton:declare ,name (,@input-type -> ,output-type))
                            (coalton:define (,name x y)
                              (coal-lib:add x y)))))
                  name))
               x
               y))))

And now, we can call the polymorphic add without wrapping!

(cl:in-package :coal-lib-user)

(coal-lib-exported:add 2 3) ;=> 5
(coal-lib-exported:add 2.0 3.0) ;=> 5.0

;; However, this errors!
(coal-lib-exported:add 2.0 3)

Unfortunately, the last of these errors. It turns out that coalton transforms the arguments into 2.0 3.0 in the last coal-lib-exported:add call above. There is another issue that above assumes the output type is the same as the first input. This will often be false. We will solve both these issues using a simple variant of coalton-impl/entry:expression-entry-point. This is defined in the attached coalton-jit-helper.lisp. (For coalton devs, the only modification is that direct-application-entry-point returns the application node instead of the lisp expression.)

Using direct-application-entry-point, we can define a function that takes in the (coalton ...) expression and returns the input and output types of the nodes. We will call this function coalton-input-output-types.

(in-package :coal-lib-exported)

(defgeneric node-input-output-types (node))

(defmethod node-input-output-types ((node coalton-impl/codegen/ast:node-locally))
  (let* ((node-let (coalton-impl/codegen/ast:node-locally-subexpr node))
         (bindings (coalton-impl/codegen/ast:node-let-bindings node-let))
         (input-types (mapcar #'coalton-impl/codegen/ast:node-type
                              (mapcar #'cdr bindings)))
         (output-type (coalton-impl/codegen/ast:node-type node-let)))
    (values input-types output-type)))

(defmethod node-input-output-types ((node coalton-impl/codegen/ast:node-direct-application))
  (let* ((rands (coalton-impl/codegen/ast:node-rands node))
         (rands (remove-if-not #'coalton-impl/codegen/ast:node-literal-p rands))
         (input-types (mapcar #'coalton-impl/codegen/ast:node-type rands))
         (output-type (coalton-impl/codegen/ast:node-type node)))
    (values input-types output-type)))

(defun coalton-input-output-types (coalton-expr)
  (let* ((*readtable* (named-readtables:ensure-readtable 'coalton:coalton))
         (string (coalton-impl/reader::print-form coalton-expr))
         (source
           (coalton-impl/source:make-source-string string
                                                   :name "<macroexpansion>")))
    (multiple-value-bind (input-types output-type)
        (node-input-output-types
         (coalton-impl/entry::direct-application-entry-point
          (with-open-stream (stream (coalton-impl/source:source-stream source))
            (coalton-impl/parser/reader:with-reader-context stream
              (coalton-impl/parser/toplevel:read-expressions stream source)))))
      (values (mapcar #'coalton-impl/typechecker/types:tycon-name input-types)
              (coalton-impl/typechecker/types:tycon-name output-type)))))

We can then rewrite out add function.

(in-package :coal-lib-exported)

(let ((cache (make-hash-table :test #'equal)))
  (defun add (x y)
    (let ((types (list (coalton-type-of x) (coalton-type-of y))))
      (funcall (alexandria:ensure-gethash
                types
                cache
                (let ((name (apply #'alexandria:symbolicate 'add
                                   (loop :for type :in types
                                         :nconcing (list '/ type)))))
                  (multiple-value-bind (input-types output-type)
                      (coalton-input-output-types `(coal-lib:add ,x ,y))
                    (let ((input-type (cons (car input-types)
                                             (loop :for type :in (cdr input-types)
                                                   :nconcing (list '* type)))))
                      (eval (print
                             `(coalton:coalton-toplevel
                                (coalton:monomorphize)
                                (coalton:declare ,name (,@input-type -> ,output-type))
                                (coalton:define (,name x y)
                                  (coal-lib:add x y)))))))
                  name))
               x
               y))))

And now, all of these work! Note that the first call will be a bit slow. But subsequent calls are going to be a lot faster! This is because, the first call is used for generating, compiling and caching the coalton function.

(cl:in-package :coal-lib-user)

(time (coal-lib-exported:add 2 3)) ;=> 5
(time (coal-lib-exported:add 2.0 3)) ;=> 5.0
(time (coal-lib-exported:add 2 3.0)) ;=> 5.0
(time (coal-lib-exported:add 2.0 3.0)) ;=> 5.0
(time (coal-lib-exported:add 2.0 3.0)) ;=> 5.0

Here are the results for the last two calls.

(coalton:coalton-toplevel
  (coalton:declare add/single-float/single-float
                   (coalton:f32 * coalton:f32 -> coalton:f32))
  (coalton:define (add/single-float/single-float x y)
    (coal-lib:add x y))) 
Evaluation took:
  0.019 seconds of real time
  0.020661 seconds of total run time (0.020661 user, 0.000000 system)
  110.53% CPU
  30 forms interpreted
  7 lambdas converted
  55,676,865 processor cycles
  3,575,680 bytes consed
  
5.0

Evaluation took:
  0.000 seconds of real time
  0.000009 seconds of total run time (0.000009 user, 0.000000 system)
  100.00% CPU
  22,088 processor cycles
  0 bytes consed
  
5.0

Note that this method is only intended for end users. Other users who want to build a library on top of your library must use coalton. This is inevitable with any other DSL too (and I think acceptable).

Converting the above to a lisp library or incorporating into coalton is left as an exercise for the reader.