Rebuilding the MS-DRG grouper in Rust

Every Medicare inpatient claim in the United States is priced by the same piece of software: the MS-DRG grouper. It reads a claim (diagnoses, procedures, age, sex, discharge status) and assigns one of roughly 770 Medicare Severity Diagnosis-Related Groups (DRGs). That group is the payment - most hospitals are paid a fixed rate per DRG, so the gap between DRG 291 (heart failure with a major complication) and DRG 293 (heart failure without one) is thousands of dollars. The difference also often depends on certain additional parameters to the grouper, like whether a certain “secondary diagnosis” counts as a “major complication” under that fiscal year's rules.

In our analytics platform, we ask grouper-shaped questions constantly. What should this claim have paid? Which denials contradict the payer's own grouping logic? What happens to the DRG if the documentation had captured that second diagnosis? Answering those questions interactively means running the grouper inside our own cloud infrastructure, and executing it billions of times on demand with slight variations of the same data. By rebuilding it in Rust with additional optimizations, we run the same grouping logic at over 100x the throughput of the original implementation.

The official CMS grouper

CMS distributes the official MS-DRG grouper as Java software. As a reference implementation, it is excellent: correct by definition, updated every October, and battle-tested by every hospital in the country.

The friction arises in deployment. The grouper requires a JVM, which averaged 464 ms of startup before the first claim groups, and roughly 74 µs per claim once the JIT has warmed up (~13.5k claims per second). That profile is fine for a nightly batch job. It is the wrong shape for predictive analytics, where a what-if slider in a dashboard may require hundreds of thousands of claims to be grouped in real time.

Conceptually, grouping is a pure function: (claim, version) → DRG. By implementing it in Rust, we can call it as a plain library: no process boundary, no warmup, no runtime dependencies, compiled to wherever we need it.

Rewriting regulatory software is exactly the kind of project that dies of a thousand paper cuts, though. The rules fill a definitions manual thousands of pages long, and they change every fiscal year. The rewrite is tractable anyway, because almost none of those rules live in code.

The grouper is mostly data

When you decompile the official grouper, you find surprisingly little logic. The actual grouping code is only around 1,500 lines. Everything else is data: an ~11 MB binary blob containing the diagnosis and procedure tables, the DRG formula lists for every MDC, cluster definitions, exclusion groups, severity attributes, and description strings.

We designed the Rust engine to parse the exact bytes the official grouper ships, extracted directly from its distribution, so every test runs against the same data the reference implementation reads. The blob format is an undocumented protobuf. Generated protobuf classes embed their own FileDescriptorProto, so a small dump tool recovers the complete .proto schema mechanically from the official jar.

// protox compiles the recovered schema; no system protoc
fn main() {
    let proto = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../proto/msdrg.proto");
    println!("cargo:rerun-if-changed={}", proto.display());
 
    let fds = protox::compile([&proto], [proto.parent().unwrap()])
        .expect("proto/msdrg.proto failed to compile");
 
    prost_build::Config::new()
        // BTreeMap for all map fields: deterministic iteration for the table façade and tests.
        .btree_map(["."])
        .compile_fds(fds)
        .expect("prost codegen failed");
}

This lets us support new fiscal years with a data drop instead of a code rewrite.

Rust features we brought to the table

The failure mode that matters in a payment engine is a wrong answer. To guard against it, we:

• Used Rust's type system to enforce the exact data model of the official grouper.
• Ensured every Java constant the wire format can express is an exhaustive Rust enum.
• Defined MDC suppression sets as u64 bitsets, a major efficiency improvement over the original implementation.
• Loaded the tables into Arcs, so we can share them between threads without copying.
• Made the iteration order of map keys deterministic, so we can write reproducible tests.
• Wrote a full differential test suite that compares against the official grouper across all eight supported versions.

Compiling a language from 1985

The most surprising thing inside the grouper is a tiny boolean formula language. There are no if statements. Each DRG is a row in a rank-ordered list, guarded by a formula string like (sensor | NORcardcath) & bypass, where each operand names a set of claim attributes. The surgical hierarchy is simply the rank order of those rows; the first formula that matches is your DRG. Evaluating a terminal can accumulate matched attributes, mark codes as used, or record a return code, so short-circuit ordering is part of the MS-DRG semantics.

After experimenting with several designs, the fastest engine parses each formula once at load time, resolves its operands against the version's tables, and compiles the AST into a flat list of ops threaded with conditional jumps. Per-claim evaluation is a single scan with a boolean accumulator: no recursion, no allocation, and exactly the reference's evaluation order.

// the core of the per-claim evaluator
pub fn evaluate<C: ClaimAttrs + ?Sized>(&self, ctx: &mut C) -> bool {
    let mut acc = false;
    let mut pc = 0usize;
    while pc < self.ops.len() {
        match self.ops[pc] {
            Op::Term(t) => acc = eval_terminal(&self.terminals[t as usize], self.mode, ctx),
            Op::JumpIfFalse(target) => if !acc { pc = target as usize; continue },
            Op::JumpIfTrue(target) => if acc { pc = target as usize; continue },
            Op::Not => acc = !acc,
        }
        pc += 1;
    }
    acc
}

Because the formulas are data, we can also prove coverage instead of hoping for it. A test compiles every formula in both committed blobs across all eight supported versions (over 10k rows), and fails if any rule produces a parse error or an operand the resolver doesn't recognize. New constructs in a future grouper version fail this test instead of slipping through.

Differential testing

The only ground truth is the official implementation itself, so we built a test harness that checks our answers against the reference grouper. Correctness means byte-for-byte equivalence on every output field.

We first built a small Java wrapper that turns the official MS-DRG grouper into a newline-delimited JSON service: one claim per line on stdin, one full grouping result per line on stdout. The test harness keeps the process alive for the whole run, amortizing the JVM startup across hundreds of thousands of claims.

The differential test runner generates claims with a seeded SplitMix64 generator, so claim i of seed s is reproducible forever. A disagreement found today can be replayed unchanged in a year. The generator is deliberately hostile: up to 24 secondary diagnoses and 24 procedures per claim, every sex and discharge status, and about 2% garbage codes, because the reference grouper's behavior on malformed input is part of its contract too:

// ~98% real codes, ~2% junk (condensed)
match rng.below(6) {
    0 => random_garbage(rng),                 // "xK9q2", both cases, length 1–9
    1 => dotted(pick(rng, pool)),             // valid code with a dot inserted
    2 => pick(rng, pool).to_ascii_lowercase(),
    3 => pick(rng, other_pool),               // procedure code in a diagnosis slot
    4 => String::new(),
    _ => "null".to_owned(),
}

When the two engines disagree, the harness greedily drops secondary diagnoses and procedures one at a time, keeping each removal only if the disagreement survives, until no single removal can preserve it. The 1-minimal counterexample goes into a failure corpus. Every fixed failure becomes a permanent regression test:

// shrink to a 1-minimal counterexample (condensed)
let mut current = claim.clone();
'outer: loop {
    for i in 0..current.sdx_len() {
        let candidate = current.without_sdx(i);
        if disagrees(sidecar, &candidate)? {
            current = candidate;
            continue 'outer;
        }
    }
    // …same pass over procedures…
    return Ok(current);
}

Before a claim is ever grouped, the parsed artifacts themselves are compared. Every formula AST our parser produces must render byte-identically to a dump of the reference parser's AST, for all eight versions. Every table lookup is cross-checked against reflective dumps from the official jar. When something is wrong, these tests point at the exact table and key instead of a mysteriously different DRG three stages later.

Where the reference implementation throws on a pathological input, our engine returns an error at the same evaluation point, so the harness can compare failure modes as well as answers.

A typical differential run is 200,000 generated claims against one version. The port advances stage by stage (preprocessing, MDC assignment, DRG selection, severity marking), and each stage lands only once its differential step runs clean. It is a slower way to write software, and the only one we'd trust for this.

Where this goes

The grouper moves billions of dollars a year, and software like that should be boring: small, fast, and provably faithful to the reference. For 21st Care, it means expected reimbursement checks, denial triage, and documentation what-ifs run anywhere our code runs, with every answer backed by an engine that is continuously cross-examined against the reference.

Questions about the engine, or want to see it against your own claims data? Talk to us.