Find the first specific enum variant in an iterator and transform it
Asked Answered
S

2

6

I have an enum with different variants and I want to find the first variant that matches then transform it by either returning the variant value or mapping it to something else.

In Scala, I'd use case classes to do something like:

data.collectFirst{ case d: DataD => d.data }

In Rust, I have to pattern match twice to achieve the same result. Is there a way to make it less verbose?

enum MyData {
    DataA(String),
    DataB(u64),
    DataC(bool),
    DataD { data: String, val: u32 },
}

fn main() {
    // test data
    let data = vec![
        MyData::DataB(42),
        MyData::DataD {
            data: "meaning of life".to_owned(),
            val: 42,
        },
        MyData::DataC(false),
    ];

    // find first one that matches and map it
    let found: Option<String> = data
        .iter()
        .find(|d| match **d {
            MyData::DataD { .. } => true,
            _ => false,
        })
        .and_then(|d| match *d {
            MyData::DataD { ref data, .. } => Some(data.to_owned()),
            _ => None,
        });
}
Spiffing answered 5/3, 2018 at 16:4 Comment(0)
D
11

Since Rust 1.30, you can use Iterator::find_map:

let found: Option<String> = data.iter().find_map(|d| match d {
    MyData::DataD { data, .. } => Some(data.to_owned()),
    _ => None,
});

Before then, you can use Iterator::filter_map and Iterator::next:

let found: Option<String> = data
    .iter()
    .filter_map(|d| match d {
        MyData::DataD { data, .. } => Some(data.to_owned()),
        _ => None,
    })
    .next();

I don't really like having big match statements in my iterator chains, so I'd normally make a method on MyData to use instead:

enum MyData {
    DataA(String),
    DataB(u64),
    DataC(bool),
    DataD { data: String, val: u32 },
}

impl MyData {
    fn as_data_d_data(&self) -> Option<&str> {
        match self {
            MyData::DataD { data, .. } => Some(data),
            _ => None,
        }
    }
}

fn main() {
    let data = vec![
        MyData::DataB(42),
        MyData::DataD {
            data: "meaning of life".to_owned(),
            val: 42,
        },
        MyData::DataC(false),
    ];

    let found: Option<String> = data
        .iter()
        .find_map(MyData::as_data_d_data)
        .map(str::to_owned);
}

In fact, a lot of my enums have this kind of pattern:

enum MyData {
    DataA(String),
    DataB(u64),
    DataC(bool),
    DataD(D), // made into a struct
}

impl MyData {
    fn is_a(&self) -> bool {
        match *self {
            MyData::DataA(..) => true,
            _ => false,
        }
    }

    fn as_a(&self) -> Option<&String> {
        match self {
            MyData::DataA(x) => Some(x),
            _ => None,
        }
    }

    fn into_a(self) -> Option<String> {
        match self {
            MyData::DataA(x) => Some(x),
            _ => None,
        }
    }

    // Repeat for all variants
}

I've even created a crate to automatically derive these functions. I've never published it because I'm pretty sure something similar already exists and I just haven't found it...

Dollarfish answered 5/3, 2018 at 16:25 Comment(3)
Instead of filter_map(...).next(), you can also use find_map(...).Ethanethane
The matches! macro can also be used to make the code more concise. Here is how to test a given variant: matches!(self, MyData::Data(_)).Choreographer
@Choreographer matches! will only help with the is_a style methods — you can't extract the data inside the enum using it. Since the main point of the question is getting the data, it doesn't seem to add much value. Additionally, in most cases I've done this, I've created a macro to generate the code for each variant. Typing out the explicit match once inside the macro wasn't a big deal.Dollarfish
F
-1
enum Foo {
    A(..)
    B(..)
}
impl Foo {
    fn equals(&self, foo: &Foo) -> bool {
        let eq = |foo: &Foo| match foo {
            A(_) => 0,
            B(_) => 1
        };
        eq(self) == eq(foo)
    }
}
Fidge answered 19/1, 2022 at 23:10 Comment(1)
Welcome to Stackoverflow. Your answer consists of only a code snippet that looks like might have something to do with the question. It would be nice if you could add some text explaining how exactly it solves the problem at hand. (By the way, your eq closure can be replaced by discriminant.)Hoax

© 2022 - 2024 — McMap. All rights reserved.