Compilation à la volée avec Ruby

Récemment, je suis tombé sur cette série d’articles d’Eli Bendersky sur la compilation à la volée (JIT compilation : Just-In-Time compilation) :

Eli utilise le C++ pour les trois premières parties, mais dans la quatrième il montre comment faire de la compilation à la volée en Python.

En lisant ça, je me suis demandé « Comment faire l’équivalent en Ruby ? ». Et bien c’est ce que nous allons voir !

Pour des raisons de simplicité et concision, les exemples de codes de cet article ne sont pas portable (limité à Linux/x86_64). Cependant, la démarche présentée est applicable sur d’autres OS/architecture.

Comment exécuter du code compilé à la volée en Ruby

Les étapes pour exécuter du code compilé à la volée sont les suivantes :

Allouer de la mémoire pour y placer le code à exécuter.
Créer une fonction à partir de l’adresse de la zone mémoire contenant le code.
Utiliser la fonction autant que l’on veut.
Libérer la mémoire utilisée.

Étant donné que nous allons devoir faire de la programmation relativement bas niveau (allocation mémoire entre autre), nous allons nous reposer sur l’excellent module ffi (dont j’ai déjà eu l’occasion de parler ici).

En première approche pour l’étape 1 on pourrait penser à FFI::MemoryPointer.new mais il y a de très fortes chances pour que cela se termine en Segmentation Fault (ou équivalent).

En effet, pour des raisons de sécurité, la mémoire est généralement allouée soit en « lecture/écriture » (RW : read/write), soit en « lecture/exécution » (RX : read/execute)¹. Et comme, de manière générale, les fonctions d’allocation (malloc & cie en C, FFI::MemoryPointer.new en Ruby, …) renvoient de la mémoire RW nous ne pourrons pas les utiliser pour allouer la mémoire où nous allons placer notre code.

Heureusement, il existe un appel système qui nous permet de spécifier les droits d’une zone mémoire lors de son allocation : mmap. Cette fonction est extrêmement pratique et a de nombreux cas d’usages, mais dans le cas présent nous allons en avoir une utilisation assez simple :

Allouer via mmap une zone mémoire de taille suffisante pour recevoir le code (cette zone sera allouée en RW).
Stocker le code dans cette zone mémoire.
Changer les permissions de la zone mémoire via mprotect pour la passer en RX.

Nous pourrions directement allouer la mémoire en RWX (et ainsi éviter d’avoir à appeler mprotect) mais ce n’est pas une bonne pratique (du point de vue sécurité) d’avoir les droits en écriture ET en exécution en même temps sur une même zone mémoire.

Wrapper mmap pour Ruby

Étant donné que mmap & cie ne sont pas disponibles de base en Ruby, contrairement à Python, nous allons faire un petit wrapper nous-même². Ce wrapper ne sera pas exposé publiquement, c’est un détail d’implémentation, et sera manipulé via une classe qui représente une fonction compilée à la volée : JitFunction.

Mais d’abord, voyons le wrapper en question.

Wrapper minimaliste autour de mmap & cie

        require 'ffi'

class JitFunction

  […]

  module Internal
    extend FFI::Library
    ffi_lib FFI::Library::LIBC

    attach_function 'mmap',     %i[pointer size_t int int int off_t], :pointer
    attach_function 'mprotect', %i[pointer size_t int], :int
    attach_function 'munmap',   %i[pointer size_t], :int

    PROT_NONE  = 0x0 # Data cannot be accessed.
    PROT_READ  = 0x1 # Data can be read.
    PROT_WRITE = 0x2 # Data can be written.
    PROT_EXEC  = 0x4 # Data can be executed.

    PROT_RW = PROT_READ | PROT_WRITE
    PROT_RX = PROT_READ | PROT_EXEC

    MAP_SHARED    = 0x1   # Share changes.
    MAP_PRIVATE   = 0x2   # Changes are private.
    MAP_ANONYMOUS = 0x20  # Don't use a file.

    MAP_FAILED = FFI::Pointer.new(:void, -1)
  end
  private_constant :Internal
end

      

Rien de bien compliqué : on déclare les trois fonctions que nous allons utiliser (mmap pour allouer la mémoire, mprotect pour changer les permissions et munmap pour libérer la mémoire) et les constantes dont nous allons avoir besoin (il existe d’autres constantes, vous pouvez consulter man 2 mmap pour la liste complète).

Attention : ce wrapper n’est pas portable (il ne fonctionnera pas pour Windows, sur d’autres UNIX les valeurs des constantes peuvent être différentes, …).

La classe JitFunction

Voyons maintenant la classe qui va nous permettre de créer et d’exécuter des fonctions compilées à la volée.

Implémentation de la classe JitFunction

        require 'ffi'

class JitFunction
  def initialize(ret, args, code)
    @size = code.bytesize
    @address = Internal.mmap(nil, @size, Internal::PROT_RW,
                             Internal::MAP_PRIVATE | Internal::MAP_ANONYMOUS,
                             -1, 0)
    raise 'cannot map memory' if @address == Internal::MAP_FAILED
    @address.put_string(0, code)
    status = Internal.mprotect(@address, @size, Internal::PROT_RX)
    raise "cannot change memory's permission" if status == -1
    @func = FFI::Function.new(ret, args, @address, convention: :default)
  end

  def call(*args)
    @func.call(*args)
  end

  def free
    Internal.munmap(@address, @size)
  end

  module Internal
    […]
  end
  private_constant :Internal
end

      

On peut voir que l’interface de la classe est composée de trois méthodes publiques :

initialize qui prend en argument le code de la fonction et qui crée un objet FFI::Function.
call qui permet d’exécuter la fonction.
free qui libère la mémoire utilisée par le code de la fonction.

Il n’y a pas grand chose à dire sur call (un simple wrapper autour de la méthode call de l’objet FFI::Function) ou free (un simple wrapper autour de munmap).

En revanche, initialize mérite que l’on s’y attarde.

La ligne 5 récupère la taille, en octet, du code de la fonction.
La ligne 6 alloue la mémoire pour le code via mmap. On remarque que la mémoire est alloué en RW (comme l’indique le paramètre PROT_RW). La mémoire allouée n’est pas accessible aux autres processus (utilisation de MAP_PRIVATE au lieu de MAP_SHARED) et n’est pas liée à un fichier sur le disque (utilisation de MAP_ANONYMOUS).
La ligne 9 lève une exception si l’allocation a échouée (Attention : en cas d’erreur mmap ne renvoie pas NULL mais une valeur bien spécifique représentée par MAP_FAILED³).
La ligne 10 copie le code de la fonction dans la zone allouée.
Étant donné que nous n’allons plus modifier la mémoire, la ligne 11 change les permissions de la zone mémoire de RW en RX et la ligne 12 lève une exception en cas d’échec.
Finalement, la ligne 13 crée un objet FFI::Function en précisant le type de retour, le type des arguments et l’adresse du code à exécuter.

Test

Nous pouvons maintenant tester la classe JitFunction en reprenant l’exemple utilisé par Eli dans son quatrième article : une fonction qui ajoute 4 à un entier passé en argument et retourne le résultat.

Exemple d’utilisation de JitFunction

        if __FILE__ == $PROGRAM_NAME
  code = "\x48\x89\xf8\x48\x83\xc0\x04\xc3"
  func = JitFunction.new(:long, [:long], code)
  begin
    puts func.call(-100)
  ensure
    func.free
  end
end

      

Lorsque l’on exécute ce script, -96 s’affiche. Ça fonctionne \o/

Le code complet de jit.rb est disponible ici.

Exemple d’utilisation : Brainfuck

Afin de montrer un exemple d’utilisation « réel » de la classe JitFunction, nous allons développer un programme qui va compiler à la volée un programme Brainfuck en assembleur x86_64 pour Linux puis l’exécuter : bf-jit.rb.

Point d’entrée

Commençons par le commencement : la fonction main.

Point d’entrée

        require 'set'

# Brainfuck's instruction set.
INSTRUCTION_SET = Set.new('[]><+-.,'.chars).freeze

if __FILE__ == $PROGRAM_NAME
  abort 'Missing argument: expected filename' if ARGV.empty?
  bf_code = File.read(ARGV[0]).chars.select! { |c| INSTRUCTION_SET.include? c }
  memory  = FFI::MemoryPointer.new(:uchar, 30_000)
  bf_prog = JitFunction.new(:void, [], compile_bf(bf_code, memory))
  begin
    bf_prog.call
  ensure
    bf_prog.free
    memory.free
  end
end

      

La ligne 7 vérifie que le script est bien appelé avec un argument (le nom du fichier source Brainfuck).
La ligne 8 lit le contenu du fichier et ne garde que les caractères qui font partie du jeu d’instructions Brainfuck.
La ligne 9 alloue la mémoire pour l’exécution du programme (pour rappel, un programme Brainfuck s’attend à avoir au moins 30 000 octets de mémoire à disposition).
La ligne 10 compile le code source et crée un objet JitFunction. Étant donné que le programme Brainfuck s’exécute sans interaction avec le reste du script on le représente par une fonction sans argument et avec un type de retour void.
Finalement, la ligne 12 exécute la fonction et les lignes 14 et 15 libèrent la mémoire allouée.

La compilation

Pour rappel, les instructions du Brainfuck sont les suivantes :

: déplace le pointeur d’une case mémoire vers la droite ;
< : déplace le pointeur d’une case mémoire vers la gauche ;
+ : incrémente la valeur stockée dans la case mémoire actuellement pointée ;
- : décrémente la valeur stockée dans la case mémoire actuellement pointée ;
. : affiche le caractère ASCII correspondant à la valeur de la case mémoire actuellement pointée ;
, : stocke la valeur ASCII du caractère lu dans la case mémoire actuellement pointée ;
[ : saute à l’instruction suivant le ] correspondant si la valeur de la case mémoire actuellement pointée est égale à zéro ;
] : saute à l’instruction suivant le [ correspondant si la valeur de la case mémoire actuellement pointée est différente de zéro.

On voit que pour interpréter du code Brainfuck nous allons avoir besoin de deux pointeurs :

un pointeur sur la mémoire : c’est le pointeur utilisé par la plupart des instructions (toutes sauf [ et ]).
un pointeur sur le code : c’est le pointeur qui permet de savoir quelle est la prochaine instruction à exécuter (il est incrémenté automatiquement après chaque instruction, sauf pour les instructions [ et ] qui peuvent le modifier de manière conditionnelle).

Dans le cadre d’un interpréteur, nous devons gérer ces deux pointeurs nous-même. En revanche, dans le cas de la compilation à la volée, nous allons générer du code machine et c’est le processeur qui va s’occuper de savoir quelle est la prochaine instruction à exécuter : nous avons seulement besoin de gérer le pointeur sur la mémoire. Pour représenter ce pointeur nous allons utiliser un registre. Cependant, on ne peut pas utiliser n’importe quel registre : il y a des conventions d’utilisation. En regardant la section Registers de Guide to x86-64 on voit que r10 est un bon choix (sa valeur est sauvegardé par l’appelant, on peut donc l’utiliser comme on le souhaite).

Avec ces informations en tête, nous pouvons esquisser le squelette de notre fonction de compilation.

Prologue et épilogue

        def compile_bf(bf, memory)
  # movabs r10, @memory (in little-endian)
  asm = "\x49\xBA".b + "#{[memory.address].pack('Q<')}"
  bf.each do |instruction|
    case instruction
      […]
    end
  end
  asm << "\xC3".b # ret
end

      

On peut voir qu’il y a du code ajouté avant (le prologue) et après (l’épilogue) les instructions générées à partir du code Brainfuck (qui vont être produites dans le bloc each).

Dans le prologue, nous initialisons le registre r10 avec l’adresse de la mémoire allouée pour l’exécution du programme Brainfuck. Deux remarques :

les processeurs x86_64 étant little-endian, nous utilisons le format Q< pour écrire l’adresse comme un entier 64-bit little-endian.
on utilise la fonction b de la classe String pour que notre chaîne de caractères soit traitées comme une suite d’octets (ce qu’elle est) et non pas comme du texte encodé (en UTF-8 par exemple).

Dans l’épilogue, nous ajoutons l’instruction ret. En effet, notre programme Brainfuck étant compilé en une fonction, il ne faut pas oublier l’instruction ret pour revenir dans le code appelant à la fin de l’exécution.

Compilation des instructions « simples »

Sur les huit instructions du Brainfuck, six peuvent être traduites directement en assembleur, indépendamment du contexte. Les deux dernières ([ et ]) demandent un peu plus d’effort et seront traitées dans la section suivante.

En reprenant la description des instructions de la section précédente et sachant que notre pointeur est stocké dans r10, nous avons les correspondances suivantes :

> : inc r10.
< : dec r10.
+ : addb [r10], 1
- : subb [r10], 1

Pour l’instruction , nous devons utiliser l’appel système read et pour . nous devons utiliser l’appel système write.

Pour faire un appel système en assembleur x86_64 sous Linux ce n’est pas très difficile : il suffit de mettre des valeurs dans des registres, puis d’utiliser l’instruction syscall. Pour savoir quelle valeur mettre dans quel registre, on peut se référer à la page Linux System Call Table for x86 64 par exemple.

Cela nous donne donc

Appel système à read

        mov rax, 0      # On veut appeler read.
mov rdi, 0      # On lit sur stdin.
mov rsi, r10    # On veut l’écrire sur la case mémoire actuellement pointée.
mov rdx, 1      # On veut lire seulement 1 octet.
syscall

      

Appel système à write

        mov rax, 1      # On veut appeler write.
mov rdi, 1      # On écrit sur stdout.
mov rsi, r10    # On veut écrire la case mémoire actuellement pointée.
mov rdx, 1      # On veut écrire seulement 1 octet.
syscall

      

Bon, nous avons maintenant l’équivalent en assembleur de nos instructions Brainfuck. Cependant, ce n’est pas suffisant car nous devons générer du code machine directement. Il nous faut donc savoir comment ces instructions sont encodées. Il y a plusieurs façons de faire cela, pour ma part j’ai utilisé rasm2 qui est un utilitaire faisant partie de radare2.

Utilisation de rasm2 pour encoder les instructions

% rasm2 -b 64 -a x86.as 'inc r10'
49ffc2
% rasm2 -b 64 -a x86.as 'dec r10'
49ffca
rasm2 -b 64 -a x86.as 'subb [r10], 1'
41802a01
% rasm2 -b 64 -a x86.as 'addb [r10], 1'
41800201

% rasm2 -b 64 -a x86.as 'mov rax, 1'
48c7c001000000
% rasm2 -b 64 -a x86.as 'mov rdi, 1'
48c7c701000000
% rasm2 -b 64 -a x86.as 'mov rsi, r10'
4c89d6
% rasm2 -b 64 -a x86.as 'mov rdx, 1'
48c7c201000000
% rasm2 -b 64 -a x86.as 'syscall'
0f05

% rasm2 -b 64 -a x86.as 'mov rax, 0'
48c7c000000000
% rasm2 -b 64 -a x86.as 'mov rdi, 0'
48c7c700000000

L’option -b indique la taille des registres et l’option -a l’architecture (une liste complète est disponible via l’option -L).

On peut maintenant compiler à la volée les instructions « simples ».

Fonction de compilation à la volée (sans support des boucles)

        # Mapping f(Brainfuck instruction) = x86_64 assembly
INSTRUCTIONS_MAPPING = {
  '>' => "\x49\xFF\xC2".b,
  '<' => "\x49\xFF\xCA".b,
  '+' => "\x41\x80\x02\x01".b,
  '-' => "\x41\x80\x2A\x01".b,
  '.' => [0x48, 0xC7, 0xC0, 0x01, 0x00, 0x00, 0x00,
          0x48, 0xC7, 0xC7, 0x01, 0x00, 0x00, 0x00,
          0x4C, 0x89, 0xD6,
          0x48, 0xC7, 0xC2, 0x01, 0x00, 0x00, 0x00,
          0x0F, 0x05].pack('c*'),
  ',' => [0x48, 0xC7, 0xC0, 0x00, 0x00, 0x00, 0x00,
          0x48, 0xC7, 0xC7, 0x00, 0x00, 0x00, 0x00,
          0x4C, 0x89, 0xD6,
          0x48, 0xC7, 0xC2, 0x01, 0x00, 0x00, 0x00,
          0x0F, 0x05].pack('c*')
}.freeze

def compile_bf(bf, memory)
  asm = "\x49\xBA".b + "#{[memory.address].pack('Q<')}"
  bf.each do |instruction|
    case instruction
    when '>', '<', '+', '-', '.', ','
      asm << INSTRUCTIONS_MAPPING.fetch(instruction)
    when '['
      […]
    when ']'
      […]
    else
      raise "invalid instruction: #{instruction}"
    end
  end
  asm << "\xC3".b
end

      

On remarque que j’ai ajouté une clause else au case. Cela semble inutile au premier abord, car bf ne contient que des caractères faisant parti du jeu d’instructions du Brainfuck (grâce au select appliqué dans la fonction main). Cependant, ce else est utile pour détecter des erreurs du développeur. Par exemple, dans ma première version du code le caractère . était présent deux fois dans la clause when de la ligne 23 et l’instruction , n’était pas géré. Grâce au else, l’erreur a été très rapidement mise en évidence (et corrigé !).

Compilation des boucles

Il ne reste plus que deux instructions à gérer : [ et ], qui permettent de faire des boucles en Brainfuck.

Pour rappel, l’instruction [ saute à l’instruction suivant le ] correspondant si la valeur de la case mémoire actuellement pointée est égale à zéro. En assembleur ça se traduit par :

Instructions assembleur pour [

        cmpb [r10], 0
jz offset_vers_l_instruction_suivant_]

Dans le cas de ] on fait un saut si la valeur est différente de zéro, ce qui donne :

Instructions assembleur pour ]

        cmpb [r10], 0
jnz offset_vers_l_instruction_suivant_[

Le problème étant de calculer les offsets des sauts (surtout dans le cas de [ car on ne connait pas encore la position du ] correspondant au moment où l’on traite [).

Pour résoudre ce problème, nous allons avoir besoin d’une pile et nous allons appliquer les opérations suivantes :

lorsque l’on rencontre l’instruction [ :
- on génère l’instruction de comparaison (le cmpb).
- on empile la position courante (jz_pos) car nous devrons y revenir pour mettre à jour le code généré lorsque l’on trouvera le ] correspondant.
- on insère un placeholder de 6 octets (6 octets car l’instruction de saut jz est encodée sur 2 octets et l’offset est encodé sur 4 octets).
lorsque l’on rencontre l’instruction ] :
- on génère l’instruction de comparaison (le cmpb).
- on récupère le jz_pos correspondant (sur le dessus de la pile) : il contient l’adresse du placeholder à modifier.
- on calcule l’offset du saut pour [ : l’offset est calculé entre les deux positions situées juste après les instructions de saut elles-mêmes.
- on remplace le placeholder par un jz avec l’offset que l’on vient de calculer.
- on calcule l’offset dans l’autre sens pour ] (pour sauter en arrière, vers le début de la boucle) et on génère l’instruction jnz avec cet offset.

Un petit schéma pour essayer de visualiser tout ça :

Schéma explicatif du calcul des offsets

Et le code correspondant :

Fonction de compilation à la volée

        def compile_bf(bf, memory)
  stack = []
  # movabs r10, @memory (in little-endian)
  asm = "\x49\xBA".b + "#{[memory.address].pack('Q<')}"
  bf.each do |instruction|
    case instruction
    when '>', '<', '+', '-', '.', ','
      asm << INSTRUCTIONS_MAPPING.fetch(instruction)
    when '['
      asm << "\x41\x80\x3A\x00".b # cmpb [r10], 0
      stack << asm.size # Save the offset where we will patch the jump.
      # Insert a placeholder of 6 bytes (2 for JZ + 4 for the address).
      asm << "\x0F\x84\x00\x00\x00\x00".b
    when ']'
      raise 'unmatched ]' if stack.empty?
      asm << "\x41\x80\x3A\x00".b # cmpb [r10], 0
      jz_pos     = stack.pop    # Address of the JZ to patch.
      start_loop = jz_pos + 6   # Address of the first instruction of the loop body.
      end_loop   = asm.size + 6 # Address of the first instruction after the loop body.
      # Generate the jump to the start of the loop (using JNZ).
      offset = compute_jump_offset(end_loop, start_loop)
      asm << "\x0F\x85".b << [offset].pack('L<') # L< because little endian!
      # Go back to patch the jump to the end of the loop.
      offset = compute_jump_offset(start_loop, end_loop)
      asm[jz_pos, 6] = "\x0F\x84".b + [offset].pack('L<')
    else
      raise "invalid instruction: #{instruction}"
    end
  end
  raise 'unmatched [' unless stack.empty?
  asm << "\xC3".b # ret
end

      

La fonction compute_jump_offset est une traduction en Ruby de la fonction C++ compute_relative_32bit_offset d’Eli qui est disponible ici. La seule subtilité c’est que les entiers en Ruby n’ont pas de taille fixe donc il faut bien penser à ne garder que les 32 premiers bits avec un masque.

La fonction compute_jump_offset

        def compute_jump_offset(from, to)
  if to >= from
    diff = to - from
    raise 'offset too large' unless diff < (1 << 31)
    diff
  else
    # Here the diff is negative, so we need to encode it as 2s complement.
    diff = from - to
    raise 'offset too large' unless (diff - 1) < (1 << 31)
    (~diff + 1) & 0xFFFFFFFF
  end
end

      

Voilà ! Notre fonction de compilation est complète. Il n’y a plus qu’a tester :).

Benchmark

Le code complet de bf-jit.rb est disponible ici.

Pour mettre en évidence les avantages de la compilation à la volée en terme de performance, j’ai testé bf-jit.rb contre deux autres implémentations parmi celles présentées dans mon premier article sur le Brainfuck :

rbfi : un interpréteur naïf écrit en Ruby, par Ouranos.
cbfi : mon interpréteur en C, qui implémente quelques optimisations basiques.

Le programme Brainfuck utilisé pour le benchmark est primes.bf qui est un programme qui calcule les nombres premiers.

Résultats

% time ./cbfi "$(cat primes.bf)" <<< 100
[…] 0,95s user 0,00s system 99% cpu 0,953 total

% time ./rbfi.rb primes.bf <<< 100
[…] 493,37s user 0,41s system 99% cpu 8:13,87 total

% time ./bf-jit.rb primes.bf <<< 100
[…] 0,23s user 0,02s system 99% cpu 0,249 total

Pas mal du tout ! Presque 4x plus rapide que la version C ! Et pourtant notre compilateur est très stupide : il traduit le code instruction par instruction.

D’ailleurs, ce comportement naïf est pénalisant sur certains programmes tel que hanoi.bf où mon interpréteur C est plus rapide (grâce à ses optimisations) que bf-jit.rb. Cela étant dit, il nous suffirait d’implémenter les optimisations décrites dans les articles d’Eli pour que notre compilateur repasse en tête :-)

Aller plus loin

Au lieu d’écrire notre code directement en hexadécimal, on peut utiliser un assembleur (comme metasm ou wilson par exemple) pour pouvoir utiliser des mnémoniques et améliorer la lisibilité/souplesse du code.

Par exemple, si l’on utilise metasm pour réécrire le petit code d’exemple de jit.rb on obtient :

Utilisation de metasm

        if __FILE__ == $PROGRAM_NAME
  require 'metasm'

  asm = <<-EOS
  mov rax, rdi
  add rax, 4
  ret
  EOS
  code = Metasm::Shellcode.assemble(Metasm::X86_64.new, asm).encoded.data
  func = JitFunction.new(:long, [:long], code)
  begin
    puts func.call(-100)
  ensure
    func.free
  end
end

      

Ce qui est tout de même un poil plus lisible :)

On pourrait même aller jusqu’à utiliser LLVM (via ruby-llvm) pour générer du code optimisé et devenir facilement multi-plateforme au lieu d’être limité aux processeurs x86_64.

Pour savoir pourquoi et comment, je vous renvoie sur l’article Wikipédia du W^X.↩
Nous aurions pu utiliser la gem mmap2, mais vu que notre cas d’utilisation est relativement simple et limité c’est plus pédagogique de le faire soi-même.↩
mmap ne peut pas renvoyer NULL pour signaler une erreur car l’adresse 0 est une valeur de retour valide (en utilisant le flag MAP_FIXED), cette fonctionnalité de mmap a d’ailleurs souvent été utilisée pour exploiter des déréférencements de pointeur NULL dans le noyau (voir cet article par exemple).↩

Compilation à la volée avec Ruby

Comment exécuter du code compilé à la volée en Ruby

Wrapper mmap pour Ruby

La classe JitFunction

Test

Exemple d’utilisation : Brainfuck

Point d’entrée

La compilation

Compilation des instructions « simples »

Compilation des boucles

Benchmark

Aller plus loin

Compilation des instructions « simples »