Implement repository with State Monad (use IO and State together) - scala

Question

I have this repository :
trait TrainRepository {
def get(trainId: TrainId): IO[Option[Train]]
def getAll: IO[List[Train]]
def save(train: Train): IO[Train]
}
I would like to provide an in-memory implementation using a State Monad.
But I am stuck :
if I extend the trait, I will be stuck with the types IO[...] so I believe that I would not be able to use a State Monad.
do I need to use a natural transformation?
do I need to use a Free Monad? (I would rather not)
How would you do that?
EDIT to give a little bit more context:
trait TrainRepository[F[_]] {
def get(trainId: TrainId): F[Option[Train]]
def save(train: Train): F[Train]
}
class TrainService[F[_]](repository: TrainRepository[F])(implicit monad: Monad[F]) {
def reservation(id: TrainId): F[Train] =
for{
train <- repository.get(id)
updatedTrain <- train match {
case None => monad.pure("test") // return error here
case Some(train) => monad.pure(train.bookSeat)
}
_ <- repository.save(updatedTrain)
} yield updatedTrain
}
type TrainStateRepository[A] = State[Map[TrainId, Train], A]
val inMemoryTrainRepository = new TrainRepository[TrainStateRepository] {
override def get(trainId: TrainId): TrainStateRepository[Option[Train]] = ???
override def save(train: Train): TrainStateRepository[Train] = ???
}
val postgresTrainRepository = new TrainRepository[IO] {
override def get(trainId: TrainId): IO[Option[Train]] = ???
override def save(train: Train): IO[Train] = ???
}
val testTrainService = new TrainService[IO](inMemoryTrainRepository)
// The error is here ^^^^
// I cannot mix IO and State
val prodTrainService = new TrainService[IO](postgresTrainRepository)

Answer 1

You can introduce a type parameter in order to abstract over your monad:
trait TrainRepository[F[_]] {
def get(trainId: TrainId): F[Option[Train]]
def getAll: F[List[Train]]
def save(train: Train): F[Train]
}
Then your implementation with state monad can look like
type TrainsState[A] = State[Map[TrainId, Train], A]
class StateTrainRepository extends TrainRepository[TrainsState] {
override def get(trainId: TrainId): TrainsState[Option[Train]] = State.inspect(_.get(trainId))
override def getAll: TrainsState[List[Train]] = State.inspect(_.values.toList)
override def save(train: Train): TrainsState[Train] =
State.modify[Map[TrainId, Train]](m => m + (train.id -> train)) *> State.pure(train)
}