Saturday, 13 August 2011

Simple Arity Checking Using Higher Order Functions

As a learning exercise I'm writing a little command line tool to connect to Remember the Milk using Clojure. The idea is that the tool is a mini shell (or repl) which reads commands, executes them, and prints the result.

The commands are written as simple Clojure functions. When I type something at the command line, the command is looked up in a map and, if found, executed, passing the rest of the arguments.

For this to be safe I need to check the arity of the function matches the number of commands entered at the command line.

For example:

When I type 'help' then all the available commands are returned:
rtm> help
(echo exit help)

If I type help with an argument, then it still works, because help is multi-arity, and has an implementation that takes an argument:
rtm> help command
Help for command

If I type help with more than one argument, the arity check kicks in and refuses to call the function:
rtm> help won't work
help: wrong number of args
Help for help

It took a bit of thinking to work out how to do this, but I eventually came up with the following solution.

;; higher order functions rock!
(defn arity-check
  "Returns a function that evaluates to true if the arity matches the count"
  [arglist]
  ;; special case - if arglist is of zero length then no need to check for & args
  (if (= 0 (count arglist))
    #(= % 0)
    (let [arg-map (apply assoc {}
                         (interleave arglist (range 0 (count arglist))))]
      ;; if & args found then number of args is >= the position of the &
      ;; otherwise it's just a simple size comparison
      (if ('& arg-map)
        #(>= % ('& arg-map))
        #(= % (count arglist))))))


;; this builds a collection of functions, one for each of the arglists
;; which evaluates to true if the number of args matches the arity of the
;; arglist. it then applies each function in turn against the size of the
;; args, and determines if any of them returned true. if at least one of
;; them returned true, then we can safely do the call
(defn arity-matches-args
  "Returns true if the args match up to the function"
  [f args]
  (let [arity-check-fns (map arity-check (:arglists (meta f)))]
    ((set (map #(% (count args)) arity-check-fns)) true)))

The arity-check function takes an arglist, as returned from the meta-data of a function. This is a higher order function which returns a function that evaluates to true if the number passed in is assignment compatible with the arity of the arglist.

e.g. if the arglist is [x] then 1 argument is expected. If it is [x y] then 2 arguments are expected. If it is [& args] then any number of arguments >= 0 are expected. If it's [x & args] then any number of arguments >= 1.

The arity-matches-args function calls the arity-check for each of the arglists in the meta data of the function, and stores them in a sequence, arity-check-fns. It then calls every one of those functions passing in the number of arguments in args, and generates a set of the results, which will be either true, false, or nil. It then uses the set as a function to see if it contains true. If it does then at least one of the arglist arities matches the number of arguments that have been entered.

There is probably an easier way to do this...but I'm still proud of this solution. It opened my eyes to the power of higher order functions

The source code is at https://github.com/NeillAlexander/rtm-clj

No comments:

Post a Comment