Resolvendo um problema de desempenho de 3 anos do Gleam com arenas em Rust

• A estrutura recursiva Document do formatador do Gleam trocou as alocações individuais no heap repetidas por referências baseadas em arena, resolvendo uma issue de melhoria de desempenho que estava aberta havia 3 anos
• Box<Self> foi substituído por &Self armazenado em uma arena, e a verificação de lifetimes do Rust garante que os dados internos não sejam referenciados depois que a arena for removida
• Ao alocar apenas uma vez e reutilizar centenas de Documents repetidos, como palavras-chave da linguagem e vírgulas, o tempo de execução do pretty printer caiu de 13 ms para 9,8 ms, ficando 24% mais rápido
• O tempo total de execução do gleam format, incluindo leitura e parsing do código-fonte, também foi reduzido em 13%, e o pico de uso de memória caiu de 8,4 MB para 7,6 MB, cerca de 10% a menos
• Foi necessário um trabalho manual amplo, de +2963/-1032, para alterar funções e chamadas a fim de passar a arena adiante, mas a abordagem de reduzir alocações individuais repetitivas melhorou tanto a velocidade quanto o uso de memória

---

Custo de alocação na estrutura recursiva Document

• O pretty printer do Gleam usa uma estrutura de dados Document recursiva que representa como imprimir o código e onde quebrar linhas quando o limite de largura é excedido
• As principais variantes incluem String, que contém o texto a ser impresso, Break, que define as strings antes e depois de uma quebra de linha, Group, que agrupa vários documentos, e Nest, que aumenta a indentação ao quebrar linhas
• Mesmo uma lista [1, 2] é representada combinando String, Break, Nest e Group
  • Se couber na linha atual, é impressa em uma única linha, como [1, 2]
  • Se exceder o limite de linha, é dividida em várias linhas conforme os Breaks definidos, com indentação, e a vírgula final também é impressa
• Variantes que contêm outro Document internamente, como Nest, antes usavam Box<Self>, portanto cada documento aninhado precisava de uma alocação individual no heap, o que podia representar uma parte significativa do tempo de execução
• O algoritmo de base do pretty printer do Gleam pode ser consultado no artigo Strictly Pretty

Substituição de Box por referências de arena

• O antigo Nest(Box<Self>) foi alterado para Nest(&'doc Self), mudando a estrutura para referenciar outro Document armazenado em uma arena
• Para armazenar referências, foi adicionado ao Document o lifetime 'doc dos dados do documento, separado do lifetime 'string das strings
• A implementação usa o crate typed_arena
  • Ao armazenar um valor na arena com alloc, uma referência para esse valor é retornada
  • Os dados alocados podem ser usados enquanto a arena estiver viva
  • Quando a arena sai de escopo, os dados internos também são removidos
• A referência retornada por alloc(&self, value: T) -> &mut T não pode viver mais do que a arena, e o borrow checker do Rust bloqueia referências inválidas após a remoção da arena

Alocando documentos repetidos uma única vez para reutilização

• A arena é usada não só para referenciar outros documentos, mas também para fazer cache de Documents repetidos em todo o código
• Não é mais necessário criar do zero, toda vez, documentos que representam palavras-chave da linguagem, como String("fn"), String("pub") e String("type")
• Documentos como Break { unbroken: ", ", broken: "," }, usados entre itens de listas, também são alocados uma vez e reutilizados
• Agora é possível alocar apenas uma vez cada uma das centenas de pequenos documentos que antes precisavam continuar sendo colocados em Box

API simples e mudanças amplas nos pontos de chamada

• A API da arena em si é simples, mas a conversão de uma grande parte do código existente exigiu alterações repetitivas
• Funções que antes apenas chamavam Box::new agora precisam receber um argumento adicional de arena onde os dados serão alocados
• Funções como format_list passaram a repassar a mesma arena também para subfunções que formatam itens da lista, e a criar documentos aninhados com arena.alloc em vez de Box::new
• Após alguns dias de buscas e substituições cuidadosas, o trabalho foi concluído em um pull request de +2963/-1032
• Foi um trabalho repetitivo e manual, e LLMs não foram usados nesse processo

Melhorias de velocidade e pico de uso de memória

• Ao formatar o projeto Gleam real squirrel, o tempo de execução do pretty printer caiu de 13 ms para 9,8 ms
  • Comparando apenas o pretty printer, isso representa uma melhoria de desempenho de 24%
• Como gleam format também lê e faz parsing do código-fonte do projeto além de formatar, o pretty printer representa apenas parte do trabalho total, mas o tempo total de execução também foi reduzido em 13%
• O pico de uso de memória caiu de 8,4 MB para 7,6 MB, uma redução de cerca de 10%
• Como resultado da redução do número de alocações individuais no heap com a arena, o formatador ficou mais rápido e o uso de memória também diminuiu

[GN⁺]

Opiniões no Lobste.rs

• Se quiser reduzir ainda mais o pico de uso de memória em 10%, vale considerar outra arena que retorne handles opacos em vez de referências. Se o handle puder ser de 16 ou 32 bits em vez de 64 bits, isso faz bastante diferença quando há muitas alocações pequenas
  Segundo https://donsz.nl/blog/arenas/, o único crate que suporta isso é o compact_arena, mas, se ele não atender aos requisitos, também não é tão difícil implementar isso por conta própria

  • Ao usar handles opacos, também dá para armazenar os dados de forma mais compacta. O projeto pessoal columnar armazena uma lista de Option<T> separando em cerca de 1 bit por item para indicar presença e os valores T dos itens que realmente existem
    Em vez de retornar &Option<T> como tipo de referência, ele retorna Option<&T>; de forma mais geral, a ideia é fazer recursão estrutural seguindo as referências, em vez de colocar referências dentro da recursão estrutural
  • O artigo foi baseado em uma pesquisa de 2 anos atrás e não cobriu perfeitamente todos os crates, então, se você encontrar erros ou bibliotecas novas que tenham ficado de fora, pode enviar um PR
    Arenas merecem uso mais amplo e, em Rust, retornar referências em vez de handles ou smart pointers torna o pattern matching em enums recursivos muito conveniente. Normalmente, algo como if let MyAst::Add(MyAst::Multiply(a, b), c) = expr {} não funciona por causa dos boxes ou ponteiros intermediários, mas com referências é possível fazer pattern matching de forma transparente sem recursos nightly. Se todas as referências apontarem para uma única arena, os lifetimes que conectam a árvore inteira também permanecem consistentes
  • Como o resultado foi bom, estão considerando migrar outras partes do compilador para o modelo de arena e pretendem continuar explorando essa abordagem

• Essa abordagem é bem parecida com a suportada por pretty, outra biblioteca de pretty-printing no estilo de Wadler. Mas, por suportar vários modos de alocação de memória, a API pode parecer um pouco pesada

• Surpreendente que o autor do artigo citado no texto trabalhe em uma empresa comum de TI em Brunswick

  • Esse é o charme do código aberto. Muitas vezes, especialização e trabalho do dia a dia não têm tanta relação entre si