Design help: Converting generics code to dyn code

Hi!

I would really appreciate code review on Rust Playground

Context

I’m working on a game-theory engine with two core traits:

  • StateEnvironment
  • AI

Multiple AIs compete within a single environment.

Each environment defines associated State and Action types. These must implement:

Eq + Hash + Clone + Debug

The AI algorithms need to store and compare states (e.g., search trees), so those bounds are required.

Previous design

pub trait StateEnvironment {
    type StateT: Eq + PartialEq + Hash + Clone + Debug + 'static;
    type ActionT: Eq + PartialEq + Hash + Clone + Debug + 'static;

    fn initial(&self) -> Self::StateT;
    fn next(&self, state: &Self::StateT, action: &Self::ActionT) -> Self::StateT;
    fn status(&self, state: &Self::StateT) -> GameStatus;
}

pub trait AI<E: StateEnvironment> {
    fn next(&mut self, state: &E::StateT, env: &E) -> E::ActionT;
}

Battle loop:

pub fn battle<'a, E: StateEnvironment>(
    e: &'a E,
    mut actors: Vec<Box<dyn Algorithm<TrackingWrapper<'a, E>>>>,
) -> Result<EndState, BattleError> {
    let mut track_env = TrackingWrapper::new(e);
    
   while let GameStatus::Playing(player_idx) = track_env.status(&state) {
        let player = &mut players[player_idx as usize];
        let action = player.next(&track_env, &state);
        state = track_env.next(&state, &action);
    }
   ....
}

The problem

The environment is often wrapped (tracking, logging, caching, sub-arena selection, etc.).

Because AI is generic over the entire StateEnvironment, any wrapper changes the environment type. That means:

  • The code constructing the AI must know the final wrapped type.
  • But the wrapper stack is often assembled inside battle().
  • This tightly couples AI construction with environment wrapping.
  • Wrappers may also change the State and Action type, so even parametrize AI by those only wouldn't really solve it

The core issue is that the AI and E types must agree on State and Action types and that's a big ol hassle, in the face of wrappers and composition. And really the AI's dont care about the types, they just need the traits.

New attempt

So I decided to completely switch over to dynamic dispatch. This way, the AI doesn't care about the specific type, it just gets something that implements the required traits. So my new main loop would look like this.

fn battle_new<'a>(
    algo1: &'a mut dyn AI,
    algo2: &'a mut dyn AI,
    env: &mut dyn GameEnv,
) {
    let mut state = env.initial();

    while let GameStatus::Playing(player_idx) = env.status(&state) {
        let player = &mut players[player_idx as usize];
        let action = player.next(env, &state);
        state = env.next(&state, &action);
    }

    println!("Result: {:?}", env.status(&state));
}

However in the effort to make this work, I feel I have created a monster

#[derive(Clone)]
pub struct ObjectRef(pub Rc<dyn ObjectAPI>);

pub trait ObjectAPI {
    fn eq_dyn(&self, other: &dyn Any) -> bool;
    fn hash_dyn(&self, state: &mut dyn Hasher);
    fn as_any(&self) -> &dyn Any;
    fn debug_dyn(&self, f: &mut Formatter<'_>) -> fmt::Result;
}

So now the AI trait can be

pub trait AI {
    fn next(&mut self, environment: &dyn GameEnv, state: &ObjectRef) -> ObjectRef;
}

Now all states and actions in my program would be wrapped in ObjectRef, killing performance, introducing a lot of boilerplate and overall forcing Rust to be more like JavaScript.

Is this how dyn is supposed to be done?

I understand that I basically welded a dynamic typing system on Rust, and wiped all of type safety with it, but... is there any way to do what I want, without sacrificing type safety? I especially don't like that Env::next() takes in an ObjectRef (wrapped State) and returns and ObjectRef (also wrapped State). I could trivially accidentally return an Action instead of a State, and the type system would be perfectly happy.

What I tried but didn't work:

battle() wouldn't get the final AI types but would get &dyn AIFactory with make<T>() -> Box<&dyn AI>

  • Failed: dyn types can't have type parameters

Skipping the ObjectRef type

  • Failed: I need the action and state types to implement Hash, Clone, Debug, Eq, and I cant do blanket impl on Rc<dyn>
1 Like

All of this is especially annoying for ActionT, which in lot of cases can even be Copy and this incurs ridiculous overhead

In situations like this one, I think it's always a good advise to pause for a moment and try to understand what these language constructs (Generics, traits, etc.) are meant for, so that we make a better use of them in our abstractions.

Traits are meant for encoding behaviour. For example:

pub trait Speak {
  fn speak(&self) -> String
}

pub struct Human {
  name: String
}

impl Speak for Human {
  fn speak(&self) -> String {
    fmt!("Hi, I'm {}", self.name)
  }
}

But traits can also be used for aggregating existing behaviour into a single trait:

pub trait ObjectAPI: Eq + Debug + Hash {}

In the example above, by stating that any implementor of ObjectAPI must also implement Eq , Hash and Debug, we can get rid of the ad-hoc methods that you introduced.

Now, traits can also be parameterized, just like functions. This means that you don't necessarily need to know what concrete type their must receive or return. You can simply let the implementor decide that.

pub trait MyParameterizedTrait<T> {
  fn do_something_with_t(t: T) -> ()
}

Finally, when talking about dyn objects in Rust, we are talking about instances of implementors of the given trait.

I hope this revisit of the basics helps you get rid of the things that care causing friction in your implementation.

1 Like

First of all thank you for looking at my code!

This means that you don't necessarily need to know what concrete type their must receive or return. You can simply let the implementor decide that.

Hmmmm, which object would be the implementor in my example? (The AI impl or the Env impl,?)

In the example above, by stating that any implementor of ObjectAPI must also implement Eq , Hash and Debug, we can get rid of the ad-hoc methods that you introduced.

No well I think i need to use the ad-hoc methods because none of those are dyn compatible. I can't say ObjectAPI: Eq.... because then I can't Rc<dyn ObjectAPI>

It would be i.e. Nim for the GameEnv trait.

Ah, I see. Then sadly that akwardness is necessary. I would use an enum instead in this case for this reason, tbh.

2 Likes

I think you did pinpiont the problem, in particular:

Can you say why this happens:

Because it feels sort of like you're asking...

AI author, build an algorithm for this game, here's the core state and actions. Oh but I might add an arbitrary combination of wrappers with new/different state and actions, so now you have to handle a combinatorial explosion of new states and actions. I might add more wrappers later too, kthxbye!

That is, it seems reasonable for an AI to be defined by a single implementation of StateEnvironment with a concrete state and actions, but there's some sort of missing layer so that the wrappers aren't really wrappers that excercise the AI, they're intrusive demands on the AI.

Yes, it doesn't happen often, but chief example is a PartialPlanEnvWrapper. This env takes another env with some action plan and creates a wrapped env in which some actions are pre-set. It must wrap the original state with extra information to properly execute the plan. Thus changing the state type.

It could it could.

The main focus of this work is however the AIs and they should be able to work with any environment. The AI algorithm must be general enough to be able to work with anything implementing the GameEnv trait. Don't think of this as Starcraft AI, think of it more like in terms of planning AI, like A*, (if you wrote an A* library , you want it to be general, not only for one problem type right)

The AIs should only ever see opaque states, actions and transitions, by design, they never could (or should!) interact with the state internals.

This should've probably been in the original description but yeah, the AIs should always always just work with states as opaque types, they should never have to inspect the state internals (and they can't even in the current generics design)

1 Like

Okay, yes, I definitely misunderstood the design space. That makes sense. Thanks for explaining.

If I think of anything useful for your design, I'll make another post. The rest of this one is a side-bar about ObjectAPI and dyn. (I'm not necessarily advocating it's use, just talking technically.)


The first point is that super-traits can ease the implementation.

-pub trait ObjectAPI {
+pub trait ObjectAPI: Any + Debug {
     fn eq_dyn(&self, other: &dyn Any) -> bool;
     fn hash_dyn(&self, state: &mut dyn Hasher);
-    fn as_any(&self) -> &dyn Any;
-    fn debug_dyn(&self, f: &mut Formatter<'_>) -> fmt::Result;
 }

 impl Debug for ObjectRef {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        self.0.debug_dyn(f)
+        self.0.fmt(f)
     }
 }

The supertrait is enough to dispatch the Debug implementation, and trait upcasting can get you as_any... sort-of. The ergonomics are worse, so you may want to keep it. But I'll come back to that subject later.


The next point is that you can dispatch to a derived Hash implementation too.

-#[derive(Eq, PartialEq, Debug)]
+#[derive(Eq, PartialEq, Debug, Hash)]
 struct NimState{ ... }

 impl ObjectAPI for NimState {
     fn hash_dyn(&self, state: &mut dyn Hasher) {
-        state.write_i8(self.matches);
-        state.write_u8(self.playing);
+        self.hash(&mut state);
     }

You may have tried this and got an error. The work around is to pass &mut &mut dyn Hasher to self.hash.


The next point is perhaps the largest point: You can't blanket implement Rc<_>, but you can implement dyn YourOwnTrait. If Trait is a local trait, dyn Trait is a local type.

// N.b. you usually want to implement for `dyn Trait + '_` when implementing
// things for `dyn Trait`, but with an `: Any` supertrait bound, our lifetimes
// are always `'static`.  (This note applies throughout this post.)

impl PartialEq<dyn ObjectAPI> for dyn ObjectAPI {
    fn eq(&self, other: &Self) -> bool {
        self.eq_dyn(other)
    }
}

impl Eq for dyn ObjectAPI {}

impl Hash for dyn ObjectAPI {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.hash_dyn(state);
    }
}

This is a good intermission point: at this stage, you could get rid of ObjectRef and use Rc<dyn ObjectAPI> instead.

The rest of this post considers what it looks like to generalize things more, and then dial them back to something more targeted. It's a lot of boilerplate up front though.


You can blanket implement ObjectAPI for everything that meets the desired bounds:

impl<T: Eq + Hash + Any + Debug> ObjectAPI for T {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn eq_dyn(&self, other: &dyn Any) -> bool {
        if let Some(other) = other.downcast_ref::<Self>() {
            self == other
        } else {
            false
        }
    }

    fn hash_dyn(&self, mut state: &mut dyn Hasher) {
        self.hash(&mut state)
    }
}

But this brings us back to as_any. If you keep that method and use this blanket implementation, things like action.0.as_any() will get you the &dyn Any for Rc<dyn ObjectAPI>, and not the erased type! You would need to be very mindful about doing something like (*action.0).as_any(). (The same would go for any other methods you used directly -- almost everything has those methods.)

A better alternative in this case is to remove as_any from the trait, but make it a method on dyn ObjectAPI only. Trait upcasting will do the right thing here:

    impl dyn ObjectAPI {
        pub fn as_any(&self) -> &dyn Any {
            self
        }
    }

However... we've also expanded this concern to, well, ~everything. But there's remedy of sorts: provide a subtrait that has to be opted in to... for each class of thing.

pub trait ActionAPI: ObjectAPI {}
pub type ActionRef = Rc<dyn ActionAPI>;
impl dyn ActionAPI {
    pub fn as_any(&self) -> &dyn ::std::any::Any {
        self
    }
}
impl PartialEq for dyn ActionAPI {
    fn eq(&self, other: &Self) -> bool {
        self.eq_dyn(other)
    }
}
impl Eq for dyn ActionAPI {}
impl Hash for dyn ActionAPI {}
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.hash_dyn(state);
    }
}

// And same for `StateAPI/StateRef`

That is the ton of upfront boilerplate. A macro would ease most of the pain. The upside is that using these elsewhere just involves deriving the normal traits, and then:

impl StateAPI for NimState {}
impl ActionAPI for NimAction {}

So, that was my adhoc tour of how to make dyn Trait with Eq and Hash a bit nicer. Again, I'm not advocating it as the solution to your design problems -- I haven't come to any firm conclusions there.

2 Likes

Thank you so much for your time and the detailed post! I really appreciate how you took time to understand the code and what I'm trying to do.

I did not know I can impl for dyn Trait, I'll try it tomorrow, it's getting quite late in my time zone. I would very much like to get rid of the ObjectRef if I can, and your way of being able to just do impl StateAPI for NimState would also be tremendously helpful.

Thanks again!!

If I understand correctly this time, the AIs only really need some sort of unique identifier... plus perhaps a way to output debug representations.

Here's an experiment using IDs instead of erased Rcs:

And various thoughts about it.

  • There's a lot of cloning. Maybe you could commit to Copy ids.

  • What started as unique States and Actions are now unique ways to convert to/from ids. This has the same runtime failure downsides as downcasting roughly. It is more work for the implementor (though depending on your ID approach, easier for some -- indices, address identity...). Might not be ideal for games with huge states not directly representable by a small number of bits.

    • (I picked an approach and went with it for NimState/NimAction, might have done it differently or changed the types themselves with more thought.)
  • Edit: There are other Debug approaches but I cleaned up my original into something palatable IMO.

Replaced debug jank notes
  • The debugging approach I took is ugly and not ideal.[1] It could be cleaned up with some effort around FormatArgs or such, maybe enough to be palatable. Alternatively it may be nicer if one could commit to one of these signatures:
    // You own all states, or like the next version with shared mutability
    fn debug_state(&self, id: &StateId) -> &dyn Debug { ... }
    
    // Write your concrete state into a dummy field and return that
    fn debug_state(&mut self, id: &StateId) -> &dyn Debug { ... }
    
    // Something custom but I might stick `&self` in it; `Box` penalty for
    // the "owned" cases which would prefer the above signatures though
    fn debug_state(&self, id: &StateId) -> Box<dyn Debug + '_>
    

  1. I'm out of practice for this sort of trick. ↩︎

Yes! Exactly!

That is really smart! Coincidentally I was looking at bytemuck - Rust recently so I was thinking about it too! But I don't think I can fit the actual non-toy states into even u128. [1]

I could conceivably make the environment store all new states in its own hash-map and just hand out indices, but that has the obvious implications of basically re-implementing pointers from scratch


[1] This work considers states as grids of at least like 10x10 tiles, with multiple units (Multi Agent Path Finding). E. g.

##########
#A..#....#
#.#.#.##.#
#.#.B.#..#
#.###.#..#
#...#.#..#
###.#.##.#
#...#..A.#
#..###..B#
########## 

Actions could be represent-able by even u8 probably (6 bits to encode unit ID, 2 bits to encode direction) + a special END_OF_TURN action (0b11111111 could be END_OF_TURN) (think of Civ-like controls if you ever played, ie you control each unit individually, and then execute end of turn)

There’s many things that we can store in the env, that don’t need to be stored in the state (they can’t change)

  • Wall positions
  • Starting unit positions + team affiliations
  • Goal positions

Each state has to encode:

  • currently_playing: Enum{0, 1}
  • Unit data[] for each unit
    • position
    • alive: bool
    • reached_goal: bool
    • team: Enum{A, B}
    • moved_this_turn: bool

Fair enough! I think bytemuck and using &[u8] would be a decent approach that doesn't require getting into type-erased pointers. Might be mildly annoying to require a no-padding struct, but I suppose there's always fallback to some sort of serialization (so long as its consistent).