Making OCaml Safe for Performance Engineering

Context and Current State of OCaml

• OCaml utilizes a uniform representation of values: every value is exactly one machine word, either an immediate (data packed in 63 bits with a marker bit) or a pointer to a heap-allocated block.
• This uniformity simplifies garbage collection (GC), as words ending in zero are always valid pointers, and words ending in one are immediates.
• Polymorphism is implemented via code specialization where a single assembly implementation handles all types because they share the same one-word representation.
• OCaml 5 introduced shared-memory parallelism with a "sane" memory model where data races are bounded in space and time, unlike C++ where undefined behavior can cause arbitrary corruption.
• The uniform representation creates performance bottlenecks for high-performance software:
  • Small data (e.g., 8-bit integers) consumes a full word, wasting memory.
  • 64-bit integers (int64) and IEEE floats require heap allocation with headers, resulting in three-word blocks and poor cache locality.
  • GC pressure increases as the collector must manage every 64-bit value and float.
• Current parallelism requires runtime synchronization checks to enforce safety; these checks are cheap but unnecessary if data races could be ruled out statically.

Design Goals for Language Extensions

• Jane Street aims to provide safe, convenient, and predictable control over performance-critical aspects without sacrificing type safety or memory safety.
• Features must be "predictable" by allowing developers to statically guarantee performance behavior rather than relying on compiler heuristics.
• Implementation follows a "pay-as-you-go" model: added complexity must be invisible to standard OCaml usage and only imposed when specific features are utilized.
• Extensions are designed to be backwards compatible with upstream OCaml, ensuring interoperability with public open-source libraries.
• The team is avoiding a language fork and plans to integrate features into mainline OCaml.

Proposed Language Features

• Narrow and Flat Data Layouts (Unboxed Types):

  • Introduces types like unboxed int32 and unboxed int64 that store raw bits without headers or pointer indirection.
  • Allows "unboxed records" to be passed as register groups rather than heap pointers, improving performance.
  • Enables arrays of unboxed records with contiguous memory layout, fixing cache locality issues.
  • Implemented via "layouts" (a form of kind) that classify types by their runtime shape (e.g., bits_8, bits_64, float_64, void, or any).
  • The garbage collector is extended to track specific fields in records to determine which values to scan, ignoring unboxed data.
  • Polymorphism is implemented by generating one copy of a function per layout rather than per specific type, a middle ground between OCaml's single copy and C++'s type-specific copies.

• Stack Allocation via Modes:

  • Addresses the expense of GC and heap allocation by allowing data to be allocated on the stack and freed via stack pointer manipulation.
  • Prevents "use-after-free" bugs by enforcing a stack discipline where values cannot escape their allocation scope (e.g., stored on the heap or returned from a function).
  • Uses "modes" (specifically local vs. global) to track value lifetimes statically without requiring explicit lifetime annotations like Rust.
  • local mode guarantees a value follows stack discipline; global is the default unconstrained legacy mode.
  • Submoding allows global values to be treated as local (restricting usage) but not vice versa.
  • Enables stack allocation for closures (e.g., anonymous functions in map) when the compiler verifies they do not escape.
  • The local mode also provides safety for I/O operations (e.g., ensuring with_file does not allow file handles to be stored on the heap).

• Data Race Freedom:

  • Extends the mode system to enforce static guarantees against data races in parallel code.
  • Introduces a "contention" dimension to track if a value has been shared across threads (uncontended vs. contended).
  • Introduces a "portability" dimension to track if a value is safe to share across threads (portable vs. non-portable).
  • Functions capturing mutable state from one thread are marked non-portable and cannot be sent to other threads.
  • The spawn API requires its argument function to be portable, while join returns values that are safe to reuse after the thread terminates.
  • Provides an API for safe shared mutable memory using phantom types to link data pointers to specific lock keys.
  • Shared memory operations (e.g., map on a shared cell) require the user to hold the specific lock key, which is non-portable and thus cannot be shared.
  • The system ensures that if memory is shared (contended), it cannot be mutated unless explicitly locked.

Implementation Status and Feedback

• Unboxed types and stack allocation are currently in production use at Jane Street; teams report easier development compared to previous workarounds.
• Data race freedom implementation is complete, and the team is in the process of annotating standard libraries with modes.
• The team is currently beta-testing the data race freedom API with internal teams before rolling it out company-wide.
• A key paper on this work, "Data Race Freedom in OCamL," won an award at POPL (Principles of Programming Languages).
• The primary challenge identified is the complexity of annotating the large existing standard library with the new modes.