Consistent mutable/immutable type patterns in Java

[illinar]

This discussion applies primarily to Java API. I haven't used or reviewed Python API. C++ API may possibly benefit as well, but since there is no GC and the type system is much more powerful the material impact is not as tangible.

WPILIB types do not follow consistent approach on types mutability. Some types (Rotation2d, Translation2d) are immutable, while other (SwerveModuleState, ChassisSpeeds) are mutable. Yet majority of methods defined on mutable types also result in new allocations, as if the types exhibit immutable behavior.

While immutability is a powerful and generally much desired trait it does come at a cost of additional memory allocations. Java is specifically at disadvantage here due to GC pressures this creates.

As an example:

public Rotation2d interpolate(Rotation2d endValue, double t) {
  return plus(endValue.minus(this).times(MathUtil.clamp(t, 0, 1)));
}

This interpolation call results in 4 allocations instead of 1. Similar situations are plentiful within the code.

I can see several possible improvements to consider in addressing this situation.

Low-hanging fruit

For immutable types representing mathematical concepts, introduce constant-like properties for typical values:

public static final Rotation2d zero = Rotation2d.fromRadians(0);
public static final Rotation2d PI = Rotation2d.fromRadians(Math.PI);

And internally "unwrap" complex calculation methods (like interpolate) to result in at most single allocation for the returned value. This should not introduce breaking changes except for situations where user's code intentionally uses referential integrity. I would argue though such situations should be extremely rare and represent very questionable practice to begin with.

Mutable counterparts

Consider introducing mutable counterparts, such as MutableRotation2d, to provide an escape hatch where a sequence of calculations needs to be performed in the user's code and only the end result is relevant. This would allow writing low-allocation code where it matters, such as on the hot code path, utilizing Builder/Freeze pattern, with the clear intent to reduce accidental abuse of mutable instance. Schematically:

Rotation2d complexCalculation(Rotation2d argument) {
  final var rotation = argument.toMutable();
  rotation.rotateBy(..);
  rotation.add(...)
  rotation.multiply(...);
  return rotation.freeze();
}

Exact naming of methods is up for discussion.

Possibly breaking changes

For types like SwerveModuleState or ChassisSpeeds decide whether to "close" them to stick to immutable pattern or keep them open and either change the semantic of existing methods or introduce new methods that mutate the instance in place. This part is more difficult to get right and more input is needed. For instance this convention can be used, which is mostly compatible with the current state of the API:

public class ChassisSpeeds {
  public static ChassisSpeeds discretize(...) {
    // constructs new instance
    return new ChassisSpeeds();
  }
  
  public void ChassisSpeeds discretize() {
    // mutates internally
  }
}

From our experience this season we had to rely on the mutability a lot in this case to retain values such as sampled SwerveDriveWheelPositions and SwerveModuleState between cycles. That said I can totally see the argument for forcing WPILIB types to be immutable and let users come up with their own internal ways to deal with persistence.

[calcmogul]

We're OK with API divergence to follow language best practices or work within language constraints, and this may be one of those times.

With that said, the mutable Java units library is proving annoying to use in the way all mutable APIs are annoying for value types. We should only look into transitioning to a mutable API after we've exhausted internal optimization opportunities.

C++ API may possibly benefit as well, but since there is no GC and the type system is much more powerful the material impact is not as tangible.

I only see this impacting Java. Mutable APIs are a compromise that trades off API usage complexity for fewer allocations caused by Java's poor memory model. Python would have negligible benefit memory-wise, and C++ would have no benefit (in fact, it would bring only downsides).

Python

Python robot code is a wrapper around dynamically allocated C++ objects. Python uses refcounting instead of a GC to manage memory, so it doesn't have to worry about accruing a bunch of garbage, then spending a lot of time later cleaning it up. I don't know what Python's allocation overhead is, but in my experience, the language's runtime performance is the bottleneck because it's an interpreted language (loops in particular are very slow).

C++

In C++, there's no downsides to using an immutable API because value types are used for everything; you get the convenience of immutability without any object management overhead. ChassisSpeeds is a POD type with a trivial destructor, so allocation and deallocation are essentially free at the -O2 optimization level. It just requires incrementing or decrementing the stack pointer, which the compiler will elide by just using the stack memory locations directly.

Value types are modern best practice in C++, which often lends itself to immutable APIs.

Java

Of the three supported languages, Java forces us to make API design compromises the most often.

In Java, allocation is expensive because usable value types don't exist. Valhalla's value types are useless because they're limited to 64 bits and aren't elided until the JIT compiler finally runs; pretty much all our objects are bigger than that.

We also have no control over when the GC runs, which means it can interrupt things like control loops for indeterminate amounts of time. Determinism is the primary reason teams should be using motor controller PID, despite collocating the controls code and main logic being easier to reason about and debug.

We also have no operator overloading, which makes any math-related code ugly and unreadable compared to its C++ and Python equivalent. We've been transitioning more of WPIMath's Java API to JNI over time because the C++ version is much faster and more maintainable, even taking JNI call overhead into account.

C# fixes a lot of these issues via methods for pausing the GC and usable value types. Moving to that is a 2027 thing though, because the C# runtime doesn't support arm32.

[ThadHouse]

It doesn't support Arm32 SoftFP, which is what the roboRIO is.

[SamCarlberg]

Valhalla's value types are useless because they're limited to 64 bits and aren't elided until the JIT compiler finally runs; pretty much all our objects are bigger than that.

I've benchmarked this with Valhalla's EA builds (based on JDK 20) and redefining the Measure class as value class Measure<U extends Unit<U>> with no other changes (including any changes to the fields). Heap allocations go to zero after escape analysis determines that the objects can be created on the stack. You're right that these optimizations only come into play after the C2 JIT runs. But if user methods that run these calculations are executed while the robot is idle (eg powered on on the field before a match), then there should be enough executions for the C2 to kick in.

I believe the 64-bit limit is only for inlining on-heap values or flattening arrays, which I'm less concerned about. Short-lived objects from intermediate calculations are the largest source of object allocations, which Valhalla can help with. Though we might have to wait another ten years for it to actually arrive 🙄

Here's a run of a simple PIDF calculation (kP * error + kI * accumulated_error + kD * error_rate + kFF) implemented using different strategies.

• "Immutable" uses ImmutableMeasure explicitly for all calculations. This gives the best chance for escape analysis to determine that the values could be inlined. However, in this benchmark, they are still heap allocated at about 3.8GB/s
• "Mutable" uses cached MutableMeasure objects to hold the result of the calculation (the calculation itself is done with + and * operations on doubles). You can see no GC happens here - which is good, because the benchmark doesn't allocate anything
• "PureValues" uses a value class ValueMeasure that has an identical API to Measure, but does not implement the interface. This is using the times and plus methods to create new ValueMeasure instances; however, the C2 JIT scalarizes the values and keeps them off the heap. There's no GC here. It's also as fast as the "Mutable" benchmarks
• "ValuesLocalOnly" uses the existing Measure API and its times and plus methods. This one has both the worst performance and the highest memory allocation rate, at 4.9GB/s

Benchmark                                                                          Mode      Cnt     Score     Error   Units
PIDFBenchmark.computePIDF_Immutable                                              sample  2203832     0.021 ±   0.001   us/op
PIDFBenchmark.computePIDF_Immutable:computePIDF_Immutable·p0.00                  sample                ≈ 0             us/op
PIDFBenchmark.computePIDF_Immutable:computePIDF_Immutable·p0.50                  sample                ≈ 0             us/op
PIDFBenchmark.computePIDF_Immutable:computePIDF_Immutable·p0.90                  sample              0.042             us/op
PIDFBenchmark.computePIDF_Immutable:computePIDF_Immutable·p0.95                  sample              0.042             us/op
PIDFBenchmark.computePIDF_Immutable:computePIDF_Immutable·p0.99                  sample              0.042             us/op
PIDFBenchmark.computePIDF_Immutable:computePIDF_Immutable·p0.999                 sample              0.167             us/op
PIDFBenchmark.computePIDF_Immutable:computePIDF_Immutable·p0.9999                sample              5.729             us/op
PIDFBenchmark.computePIDF_Immutable:computePIDF_Immutable·p1.00                  sample            660.480             us/op
PIDFBenchmark.computePIDF_Immutable:·gc.alloc.rate                               sample        9  3823.222 ± 596.207  MB/sec
PIDFBenchmark.computePIDF_Immutable:·gc.alloc.rate.norm                          sample        9    32.000 ±   0.001    B/op
PIDFBenchmark.computePIDF_Immutable:·gc.count                                    sample        9  1435.000            counts
PIDFBenchmark.computePIDF_Immutable:·gc.time                                     sample        9   653.000                ms

PIDFBenchmark.computePIDF_Mutable                                                sample  2251207     0.023 ±   0.001   us/op
PIDFBenchmark.computePIDF_Mutable:computePIDF_Mutable·p0.00                      sample                ≈ 0             us/op
PIDFBenchmark.computePIDF_Mutable:computePIDF_Mutable·p0.50                      sample                ≈ 0             us/op
PIDFBenchmark.computePIDF_Mutable:computePIDF_Mutable·p0.90                      sample              0.042             us/op
PIDFBenchmark.computePIDF_Mutable:computePIDF_Mutable·p0.95                      sample              0.042             us/op
PIDFBenchmark.computePIDF_Mutable:computePIDF_Mutable·p0.99                      sample              0.042             us/op
PIDFBenchmark.computePIDF_Mutable:computePIDF_Mutable·p0.999                     sample              0.166             us/op
PIDFBenchmark.computePIDF_Mutable:computePIDF_Mutable·p0.9999                    sample              9.068             us/op
PIDFBenchmark.computePIDF_Mutable:computePIDF_Mutable·p1.00                      sample            144.384             us/op
PIDFBenchmark.computePIDF_Mutable:·gc.alloc.rate                                 sample        9     0.024 ±   0.005  MB/sec
PIDFBenchmark.computePIDF_Mutable:·gc.alloc.rate.norm                            sample        9    ≈ 10⁻⁴              B/op
PIDFBenchmark.computePIDF_Mutable:·gc.count                                      sample        9       ≈ 0            counts

PIDFBenchmark.computePIDF_PureValues                                             sample  2026871     0.025 ±   0.002   us/op
PIDFBenchmark.computePIDF_PureValues:computePIDF_PureValues·p0.00                sample                ≈ 0             us/op
PIDFBenchmark.computePIDF_PureValues:computePIDF_PureValues·p0.50                sample                ≈ 0             us/op
PIDFBenchmark.computePIDF_PureValues:computePIDF_PureValues·p0.90                sample              0.042             us/op
PIDFBenchmark.computePIDF_PureValues:computePIDF_PureValues·p0.95                sample              0.042             us/op
PIDFBenchmark.computePIDF_PureValues:computePIDF_PureValues·p0.99                sample              0.042             us/op
PIDFBenchmark.computePIDF_PureValues:computePIDF_PureValues·p0.999               sample              0.166             us/op
PIDFBenchmark.computePIDF_PureValues:computePIDF_PureValues·p0.9999              sample             10.010             us/op
PIDFBenchmark.computePIDF_PureValues:computePIDF_PureValues·p1.00                sample            665.600             us/op
PIDFBenchmark.computePIDF_PureValues:·gc.alloc.rate                              sample        9     0.026 ±   0.012  MB/sec
PIDFBenchmark.computePIDF_PureValues:·gc.alloc.rate.norm                         sample        9    ≈ 10⁻⁴              B/op
PIDFBenchmark.computePIDF_PureValues:·gc.count                                   sample        9       ≈ 0            counts

PIDFBenchmark.computePIDF_ValuesLocalOnly                                        sample  2018818     0.052 ±   0.005   us/op
PIDFBenchmark.computePIDF_ValuesLocalOnly:computePIDF_ValuesLocalOnly·p0.00      sample                ≈ 0             us/op
PIDFBenchmark.computePIDF_ValuesLocalOnly:computePIDF_ValuesLocalOnly·p0.50      sample              0.041             us/op
PIDFBenchmark.computePIDF_ValuesLocalOnly:computePIDF_ValuesLocalOnly·p0.90      sample              0.042             us/op
PIDFBenchmark.computePIDF_ValuesLocalOnly:computePIDF_ValuesLocalOnly·p0.95      sample              0.083             us/op
PIDFBenchmark.computePIDF_ValuesLocalOnly:computePIDF_ValuesLocalOnly·p0.99      sample              0.084             us/op
PIDFBenchmark.computePIDF_ValuesLocalOnly:computePIDF_ValuesLocalOnly·p0.999     sample              0.875             us/op
PIDFBenchmark.computePIDF_ValuesLocalOnly:computePIDF_ValuesLocalOnly·p0.9999    sample             30.528             us/op
PIDFBenchmark.computePIDF_ValuesLocalOnly:computePIDF_ValuesLocalOnly·p1.00      sample            998.400             us/op
PIDFBenchmark.computePIDF_ValuesLocalOnly:·gc.alloc.rate                         sample        9  4971.241 ± 709.990  MB/sec
PIDFBenchmark.computePIDF_ValuesLocalOnly:·gc.alloc.rate.norm                    sample        9   128.001 ±   0.001    B/op
PIDFBenchmark.computePIDF_ValuesLocalOnly:·gc.count                              sample        9  1670.000            counts
PIDFBenchmark.computePIDF_ValuesLocalOnly:·gc.time                               sample        9  1010.000                ms

[calcmogul]

I believe the 64-bit limit is only for inlining on-heap values or flattening arrays, which I'm less concerned about.

I'm concerned. That's a lot of heap allocation for geometry and kinematics objects the GC will eventually have to take care of. Profiling and team experience has already shown geometry/kinematics class heap allocation is an issue, and the 64-bit value type limit means we can't eliminate the allocation without rewriting the classes to use a worse mutable API.

[SamCarlberg]

I've done more digging, and profiling with a JDK build from the latest Valhalla HEAD instead of the 18-month-old early access build they have available for download.

The 64-bit limit seems to only come into play for non-final, non-volatile fields of a value type in a non-value class. A value class's fields are final, and value objects embedded inside other value objects should be able to be flattened. However, the current build doesn't support this - even for embedding a minimal value object that wraps a single byte - so Valhalla is clearly still a work in progress.

There doesn't seem to be a limit on the size of a value class for instances to be scalarized and avoid heap allocations. I tested up to 640 bits (ten doubles) without any heap allocations.

I expect intermediary geometry and kinematics objects to be scalarized once Valhalla adds support for flattening nested value objects.

[SamCarlberg]

Also see: https://mail.openjdk.org/pipermail/valhalla-spec-observers/2023-May/002243.html

Flattening

Flattening in the heap is an emergent property, which is achieved when we give up the degrees of freedom that would prevent flattening:

• Identity prevents flattening entirely;
• Nullability prevents flattening in the absence of heroics involving exotic representations for null;
• The inability to use a class without initialization requires nullability at the VM representation level, undermining flattening;
• Atomicity prevents flattening for larger value objects.

Putting this together, the null-exclusion type of implicitly constructible value classes is flattenable in the heap when the class is non-atomic or the layout is suitably small. For ordinary value classes, we can still get flattening in the calling convention: all identity-free types can be flattened on the stack, regardless of layout size or nullability.

(emphasis mine)