Swift 3.1 – prefix(while:) and drop(while:)

Just a hot minute too late for Swift 3.0, these two super-useful sequence methods have been implemented for Swift 3.1, and can be used right now on the development branch.

Let's take a look at what we can do with these functions. Suppose you have a list of messages in the chat section of your app, and you want to know how many messages users send before receiving a reply. You can think of each chat's list of messages in "chunks", with each chunk consisting of consecutive messages sent by a single user. How can we use prefix(while:) and drop(while:) to turn a single list of messages into these chunks?

(Spoiler: If you want an even nicer way to use these methods together, check out my follow-up post).

typealias User = String

struct Message {
  let user: User
  let text: String
}

func groupByUser(_ messages: [Message]) -> [[Message]] {
  guard let firstMessage = messages.first, messages.count > 1 else {
    return [messages]
  }
  
  let sameUserTest: (Message) -> Bool = {
    $0.user == firstMessage.user
  }
  let firstGroup = Array(messages.prefix(while: sameUserTest))
  let rest = Array(messages.drop(while: sameUserTest))
  
  return [Array(firstGroup)] + groupByUser(Array(rest))
}

So we find the first message, and use prefix(while:) to get all of the consecutive messages at the beginning of the array that have the same user as the first message. We then use drop(while:) to drop those same elements, and return to us the rest of the list, which we pass back in to our function to continue grouping.

We can use it like this:

let messages = [
  Message(user: "Achilles", text: "Hello? Hello?"),
  Message(user: "Achilles", text: "How do you turn this thing on?"),
  
  Message(user: "Tortoise", text: "I'm not quite sure, myself."),
  
  Message(user: "Achilles", text: "Oh, Mr. T! What a nice surprise."),
  Message(user: "Achilles", text: "I'm just trying out our dear friend the Crab's new message-transmission device."),
  Message(user: "Achilles", text: "It can decode and display any kind of message, you know."),
  
  Message(user: "Tortoise", text: "Yes, I dropped by Mr. Crab's house earlier and picked up the companion device."),
  Message(user: "Tortoise", text: "I'm quite excited to try it out – I have a specially-encoded message for just this occasion. Here it comes...")
]

let grouped = groupByUser(messages)
grouped.forEach {
  print("")
  $0.forEach {
    print($0, terminator: "\n")
  }
}

Which results in the following output:

Message(user: "Achilles", text: "Hello? Hello?")
Message(user: "Achilles", text: "How do you turn this thing on?")

Message(user: "Tortoise", text: "I\'m not quite sure, myself.")

Message(user: "Achilles", text: "Oh, Mr. T! What a nice surprise.")
Message(user: "Achilles", text: "I\'m just trying out our dear friend the Crab\'s new message-transmission device.")
Message(user: "Achilles", text: "It can decode and display any kind of message, you know.")

Message(user: "Tortoise", text: "Yes, I dropped by Mr. Crab\'s house earlier and picked up the companion device.")
Message(user: "Tortoise", text: "I\'m quite excited to try it out – I have a specially-encoded message for just this occasion. Here it comes...")

Now we have the messages grouped into array "chunks" of consecutive messages by the same author. A quick map and reduce and we have the average length of these chunks:

let averageConsecutiveMessages = Float(grouped.map { $0.count }.reduce(0, +)) / Float(grouped.count)
// 1.6

Great, we've made use of these new sequence methods to get what we were looking for. Now, I wonder what the Tortoise's last message was…