Cómo resolver con arenas de Rust un problema de rendimiento de 3 años en Gleam

• Al cambiar las asignaciones individuales en el heap que se repetían en la estructura recursiva Document del formateador de Gleam por referencias basadas en arena, se resolvió un problema pendiente de mejora de rendimiento que llevaba 3 años abierto
• Se reemplazó Box<Self> por &Self almacenado en una arena, y con la verificación de lifetimes de Rust se garantiza que los datos internos no puedan referenciarse después de que la arena haya sido eliminada
• Al asignar una sola vez y reutilizar cientos de Document repetidos, como palabras clave del lenguaje y comas, el tiempo de ejecución del pretty printer bajó de 13ms a 9.8ms, una mejora del 24%
• El tiempo total de ejecución de gleam format, incluyendo lectura del código fuente y parsing, también se redujo en 13%, y el uso máximo de memoria bajó de 8.4MB a 7.6MB, cerca de un 10%
• Aunque fue necesario modificar ampliamente funciones y sitios de llamada para pasar la arena, con un trabajo manual de +2963/-1032, fue posible mejorar tanto la velocidad como el uso de memoria reduciendo las asignaciones individuales repetitivas

---

El costo de asignación en la estructura recursiva Document

• El pretty printer de Gleam usa una estructura de datos recursiva llamada Document para representar cómo debe imprimirse el código y dónde dividirlo cuando se supera el límite de línea
• Sus variantes principales incluyen String, que contiene el texto a imprimir; Break, que define las cadenas antes y después del salto de línea; Group, que agrupa varios documentos; y Nest, que aumenta la indentación al hacer un salto de línea
• Incluso una lista como [1, 2] se representa combinando String, Break, Nest y Group
  • Si cabe en la línea actual, se imprime en una sola línea como [1, 2]
  • Si supera el límite de línea, se divide en varias líneas según el Break especificado, con indentación y mostrando también la coma final
• Las variantes que contienen otro Document internamente, como Nest, usaban antes Box<Self>, por lo que cada documento anidado requería una asignación individual en el heap, y ese trabajo podía ocupar una parte importante del tiempo de ejecución
• El algoritmo base del pretty printer de Gleam puede consultarse en el paper Strictly Pretty

Reemplazar Box por referencias de arena

• Se cambió la estructura de Nest(Box<Self>) a Nest(&'doc Self) para que referencie otro Document almacenado en la arena
• Para poder guardar referencias, se agregó a Document el lifetime 'doc de los datos del documento, separado del lifetime 'string de las cadenas
• Para la implementación se utilizó el crate typed_arena
  • Al guardar un valor en la arena con alloc, se devuelve una referencia a ese valor
  • Los datos asignados pueden usarse mientras la arena siga viva
  • Cuando la arena sale de alcance, los datos internos también se eliminan junto con ella
• La referencia devuelta por alloc(&self, value: T) -> &mut T no puede vivir más que la arena, y el borrow checker de Rust bloquea referencias inválidas después de que esta se destruye

Asignar una vez y reutilizar documentos repetidos

• La arena no solo se usó para referenciar otros documentos, sino también para cachear Document repetidos en todo el código
• Ya no hace falta crear una y otra vez documentos como String("fn"), String("pub") o String("type") para representar palabras clave del lenguaje
• Documentos como Break { unbroken: ", ", broken: "," }, usados entre elementos de una lista, también se asignan una sola vez y luego se reutilizan
• Así, cientos de documentos pequeños que antes había que meter continuamente en Box ahora pueden asignarse una sola vez cada uno

Una API simple, pero cambios extensos en los sitios de llamada

• Aunque la API de la arena en sí es simple, el proceso de migrar una gran parte del código existente requirió cambios repetitivos
• Las funciones que antes solo llamaban a Box::new ahora también necesitan recibir un argumento de arena donde asignar los datos
• Funciones como format_list ahora pasan la misma arena a las funciones internas que formatean los elementos de la lista, y crean documentos anidados con arena.alloc en lugar de Box::new
• Tras varios días de buscar y reemplazar con cuidado, el trabajo se completó en un pull request de +2963/-1032
• Fue un trabajo repetitivo y manual, y en este proceso no se utilizó ningún LLM

Mejora en velocidad y uso máximo de memoria

• Al formatear squirrel, un proyecto real de Gleam, el tiempo de ejecución del pretty printer bajó de 13ms a 9.8ms
  • Comparando solo el pretty printer, eso representa una mejora de rendimiento del 24%
• Como gleam format no solo formatea, sino que también lee y parsea el código fuente del proyecto, el pretty printer ocupa solo una parte del trabajo total, pero aun así el tiempo total de ejecución también se redujo en 13%
• El uso máximo de memoria bajó de 8.4MB a 7.6MB, aproximadamente un 10%
• Al reducir el número de asignaciones individuales en el heap mediante arenas, se mejoraron tanto la velocidad de ejecución del formateador como el uso de memoria

[GN⁺]

Opiniones en Lobste.rs

• Si se quiere reducir otro 10% el uso máximo de memoria, valdría la pena considerar otra arena que devuelva handles opacos en lugar de referencias. Los handles pueden ser de 16 o 32 bits en vez de 64 bits, lo que tiene un gran efecto cuando hay muchas asignaciones pequeñas.
  Según https://donsz.nl/blog/arenas/, el único crate que soporta esto es compact_arena, pero si no se ajusta a las necesidades, tampoco es difícil implementarlo directamente.

  • Usar handles opacos también permite almacenar los datos de forma más compacta. columnar, un proyecto personal, almacena listas de Option<T> separándolas en poco más de 1 bit por elemento para indicar presencia y los T de los elementos que realmente existen.
    A cambio, devuelve Option<&T> como tipo de referencia, no &Option<T>; de forma más general, es un enfoque de recursión estructural siguiendo referencias, en lugar de poner referencias dentro de la recursión estructural.
  • El artículo se basó en una investigación de hace 2 años y no cubre todos los crates a la perfección, así que si alguien encuentra errores o nuevas bibliotecas faltantes, puede enviar un PR.
    Las arenas merecen usarse más ampliamente, y devolver referencias en lugar de handles o smart pointers hace que el pattern matching de enums recursivos en Rust sea muy cómodo. Normalmente, aunque uno quiera escribir algo como if let MyAst::Add(MyAst::Multiply(a, b), c) = expr {}, no es posible por las cajas o punteros intermedios, pero las referencias se pueden usar de forma transparente en pattern matching sin funciones de nightly. Si todas las referencias apuntan a una misma arena, la vida útil que conecta todo el árbol también se mantiene de forma consistente.
  • Como los resultados fueron buenos, se está considerando migrar otras partes del compilador al enfoque de arenas, y se seguirá explorando este enfoque.

• Este método se parece bastante al enfoque que soporta pretty, otra biblioteca de pretty printing al estilo Wadler. Sin embargo, como soporta varios métodos de asignación de memoria, la API puede sentirse algo pesada.

• Sorprende que el autor del paper citado en el artículo trabaje en una empresa de TI común y corriente en Brunswick.

  • Esa es la gracia del open source. La especialización y el trabajo de sustento muchas veces no tienen mucha relación entre sí.