Cache intermediate results in const-generic function (using interior mutability)

I'm working on a performance critical section of my code. I want to cache some intermediate results that only rely on the size of the input, not the input itself. I'm using a const OnceLock for the caching. I'm running into a clippy warning that I'm not sure is fine to suppress.

The function in question is an internal helper function. I control all the call sites of the function. The part relevant to caching can be summarized as:

fn compute_something<const INPUT_LEN: usize>(input: _) {
    const FOO: OnceLock<_>= OnceLock::new();
    let foo = FOO.get_or_init(|| /* uses INPUT_LEN */);
    /* rest of the computation */
}

The relevant warning is declare_interior_mutable_const. In its help section, two solutions are proposed: using a static item over the const, or turning the (surrounding?) function into a const fn in order to communicate to the call site that a new const is being created. statics don't do what I want, since I want one const per call site, not one static shared across monomorphizations of the function (playground):

use std::sync::OnceLock;

fn main() {
    assert_eq!(0, use_static::<0>());
    assert_eq!(1, use_static::<1>()); // assertion failed; left: 1, right: 0
}

fn use_static<const X: usize>() -> usize {
    static FOO: OnceLock<usize>= OnceLock::new();
    *FOO.get_or_init(|| /* crunch numbers*/ X)
}

I'm not sure what the second part of the advice means. Should I turn my compute_something function into a const fn? That might be possible somehow, but other things are happening in there, and those things don't need to be const (for any other reason). Is the advice to factor out the const FOO into a new const fn? If so, how does that change anything?

The lint's description lists some cases in which using a const with interior mutability is fine, but I'm not sure any of those apply to my case. I'm using the latest stable version, so “prior to const fn stabilization” does not apply. The other exception, “types which contain private fields with interior mutability”, seemingly also does not apply, as the const is not a field on a type. Is my case spiritually identical to the latter exception?

Is there a “more correct” approach to this kind of caching than what I'm currently doing?

your code does not do what you want, const items inside a function is NOT parameterized by the enclosing function's generic parameters. for example, the following code will NOT compile:

fn zeros<const N: usize>() -> [u8; N] {
    // error[E0401]: can't use generic parameters from outer item
    const value: [u8; N] = [0; N];
    value
}

what's more importantly, even if you can work around the const limitation with associated items, OnceLock will NOT work as const, only static will work.

parameterizing OnceLock by a const generic will not work either way.

It indeed is not. The caching is entirely ineffective. Every time a const is referenced, a copy of it is made and then this copy is modified. As such each call to your function will start with an empty "cache".

1 Like

The only way to get this effect is to call a macro, not a function, that declares a static each time it is called/expanded.

Could you point me to an example or existing crate for that? I tried using a declarative macro, but I'm not sure how to increase the number of sites it's expanded at.

fn main() {
    assert_eq!(0, use_static::<0>());
    assert_eq!(1, use_static::<1>()); // assertion failed; left: 1, right: 0
}

fn use_static<const X: usize>() -> usize {
    cache!(X)
}

macro_rules! cache {
    ($x:ident) => {{
        static CACHE: OnceLock<usize> = OnceLock::new();
        *CACHE.get_or_init(|| /* crunch numbers */ $x)
    }};
}

(playground)

I suppose I can define a bunch of fn use_static_i() { cache!(i) }, but I don't think that's what you mean, because then I don't really need the macro. Perhaps a procedural macro that uses the value of X in the identifiers of the statics it creates?

You said you control all the call sites which you want to have separate caches. I’m saying that, in order to do this, you have to make those call sites into macro call sites instead of function call sites (or do something else equivalent to that). That is, your use_static must be a macro, not a function.

1 Like

Oh, I see. That makes sense. Thanks. :slightly_smiling_face: