Polymorphic methods in enums

Since I started using Swift, I've been drawn to enums for their features (associated values, type safety, etc.)

However, I noticed that in doing so, I lost some polymorphism. Consider the traffic light enum:

enum TrafficLight {
    case green
    case yellow
    case red
}

Let's say I want to add a canAdvance function, I'll end up with something like this (discussing if advancing in yellow is fine is out of the scope ;) )

enum TrafficLight {
    case green
    case yellow
    case red

    func canAdvance() -> Bool {
        switch self {
            case .green:
                return true
            case .yellow, .red:
                return false
        }
    }
}

This works, but if your enum grows and needs more methods, we end up with a switch within each method.

What I'd love to see is a way to remove those switches, that can enable tidier code. I'm thinking on something like this:

enum TrafficLight {
    case green
    case yellow
    case red

    func canAdvance() -> Bool {
        return false
    }
}

extension TrafficLight where case .green {
    func canAdvance() -> Bool {
        return true
    }
}

This would enable extensions on enums for each different case.

I don't know if this is easy (or even possible) to implement, plus, this is my first time writing on the forum, so any comments are welcome.

3 Likes
extension TrafficLight  {
    func canAdvance() -> Bool {
        if case .green = self {
            return true
        } else {
            return false
        }
    }
}
1 Like

Sounds like constant expressions in generic where clause.

Also what about associated values?

I'm not sure how would I manage the syntax, but I guess we can access the associated value in the extension.

Maybe declaring the variable name in the same where clause, just as you do in if case.

This proposal doesn't seem related to the original post. Compile time constant expressions, as far as I can tell, don't do much for extension parameters. That comes closer to having values at the type level, which has been discussed, but is not that proposal.

Swift already provides polymorphism with classes and protocols. I don't know if enums will eventually enter this dance, too.

What I know is that you're not the only one who faces this dilemma: as a code base grows, the limitations of initial design choices start to appear. Swift has a few traps like that. Strengthening one's intuition on this matter is an acquired skill. Discussing those issues is important so that the community gets a large feedback on eventual language caveats. Finding which ones should be solved, and how, is another question ;-)

For the record, refactoring your particular case around protocols could give something along:

// What is a TrafficLight?
protocol TrafficLight {
    // customization point
    var canAdvance: Bool { get }
}
extension TrafficLight {
    // default value
    var canAdvance: Bool { return false }
}

// Three particular traffic lights
struct GreenLight: TrafficLight { var canAdvance: Bool { return true } }
struct OrangeLight: TrafficLight {}
struct RedLight: TrafficLight {}

// Using traffic lights
func handle(_ trafficLight: TrafficLight) {
    if trafficLight.canAdvance {
        print("drive")
    } else {
        print("wait")
    }
}

handle(GreenLight())  // drive
handle(OrangeLight()) // wait

This refactoring is how I usually handle it. But given that Swift affords us great enums, extending them to support this seemed like a good idea.

1 Like

This kind of feature reminds me of functional languages that allow you to provide different implementations depending on the runtime value of some parameter.

For example elixir:

  def changeset(%User{} = user, %{"password" => pw} = attrs) when pw != "" do
    attrs = %{attrs | "password" => Bcrypt.hashpwsalt(pw)}

    user
    |> cast(attrs, [:firstname, :lastname, :username, :password])
    |> validate_required([:firstname, :lastname, :username, :password])
    |> unique_constraint(:username)
  end

  def changeset(%User{} = user, attrs) do
    user
    |> cast(attrs, [:firstname, :lastname, :username, :password])
    |> validate_required([:firstname, :lastname, :username, :password])
    |> unique_constraint(:username)
  end

or Haskell

factorial 0 = 1
factorial n = n * factorial (n - 1)
1 Like

Those features are almost always just sugar for a switch, but I agree that they're awfully convenient anyway.

3 Likes

I wonder if this kind of feature would be feasible to add given the existing space engine for exhaustiveness checking if all this would be is a sugar for a single implementation with a switch deciding the body.

Most of the JVM languages, Java/Scala/Kotlin, allow this. It is a feature I mis from these languages. In these languages each case of the enum is a seperate type with its own dispatch table. You can mimic this behaviour in Swift with a protocol for the enum and a struct implementing the protocol for each case but then you can’t then stop people adding new cases.

The Java syntax would probably work in Swift:

enum TrafficLight {
    case green {
        func canAdvance() -> Bool { return true }
    }
    case yellow {
        func canAdvance() -> Bool { return true }
    }
    case red {
        func canAdvance() -> Bool { return false }
    }

    func canAdvance() -> Bool // Same syntax as declaring a func in a protocol. 
}
1 Like

See, this bit I don't get. That's more verbose than using a switch.

EDIT: I too think the current language feature is a little clunky. I just haven't seen anything less clunky yet.

14 Likes

I think it depends on the number of cases the enum has, and the number of methods in there.

I agree that the syntax plays a big role in how readable is.

It also doesn't work with unlabeled associated values, unfortunately. How would I refer to the integer in this example:

enum Foo {
  case bar(Int) {
    func someFunction() -> Int {
      return ???
    }
  }
}

We'd have to do something like graft in closure shorthand variables ($0?) or invent some other syntax for it.

But I agree with @jrose; every time I've thought of this problem, I haven't been able to figure out a sweet spot that doesn't just feel like shuffling the deck chairs while providing a concrete benefit.

This could be solved with a generalization of this to all functions.

enum Color {
  case red
  case green
  case blue
  case other(Int, Int, Int)

  func something() when self = .other(50, g, b) {
  }

  func something() when self = .other(r, g, b) {
  }

  func something() { } // everything else
}

Agreed about none of the alternatives seeming like improvements. Haskell's syntax for this is elegant, but it works because (1) the type signature is written separately and (2) the switched-on parameters are all explicit and so can be easily pattern-ized. In other words, it works because there's already some verbosity that people take for granted.

I wonder how much of the clunkiness here is just the extra level of bracing required. Like if you had:

func canAdvance() -> Int = switch { // implicitly over self
  case .green: return true
  default: return false
}
8 Likes

From my experience, this would help a bit. Most enum methods I've written are generally very short, and very commonly a mapping between things. The biggest worry I have is it would erode some notions that statements in Swift aren't expressions. It begins to look like syntax you'd find in languages where this is the case.

1 Like

Another riff on this general idea, maybe we could allow top-level case labels in function bodies:

func foo(x: A, y: B) -> String? {
case (x: .foo, y: .bar):
  return "foo bar"
case (x: .zim, y: .zang):
  return "zim zang"
default:
  return nil
}
6 Likes

Hmm. Is that implicitly switching over all the non-self arguments packaged up as a tuple?

Implicitly switching over just self makes sense to me in an instance method. That's very analogous to the sorts of sugar used in other languages.

2 Likes

Sure, maybe only implicitly switching over self makes more sense. If there's an elegant syntax to allow matching either self or other arguments or some combination thereof that might be nice.