Just like partial functions in mathematics, a partial function in Scala is a function whose domain doesn’t cover all elements of the domain’s data type. For example:

val f: Function[Int, Int] = x => 100 / x

f(1)
// res1: Int = 100

f(2)
// res2: Int = 50

f(0)
// java.lang.ArithmeticException: / by zero ...

It’s a function defined for all non-zero integers, but f(0) would produce a `java.lang.ArithmeticException`.

By defining it as a partial function like below:

val pf: PartialFunction[Int, Int] = { case x if x != 0 => 100 / x }
// pf: PartialFunction[Int,Int] = 

we will be able to leverage PartialFunction’s methods like isDefinedAt to check on a given element before applying the function to it.

pf.isDefinedAt(1)
// res1: Boolean = true

pf.isDefinedAt(0)
// res2: Boolean = false

Methods lift and unlift

Scala provides a method `lift` for “lifting” a partial function into a total function that returns an Option type. Using the above partial function as an example:

val pf: PartialFunction[Int, Int] = { case x if x != 0 => 100 / x }

val f = pf.lift
// f: Int => Option[Int] = 

f(1)
// res1: Option[Int] = Some(100)

f(0)
// res2: Option[Int] = None

Simple enough. Conversely, an Option-typed total function can be “unlifted” to a partial function. Applying `unlift` to the above function `f` would create a new partial function same as `pf`:

val pf2 = f.unlift
// pf2: PartialFunction[Int,Int] = 

pf2.isDefinedAt(1)
// res3: Boolean = true

pf2.isDefinedAt(0)
// res4: Boolean = false

Function compositions

For simplicity, we’ll look at only functions with arity 1 (i.e. `Function1`, which takes a single argument). It’s trivial to use the same concept to apply to `FunctionN`.

Methods like `andThen` and `compose` enable compositions of Scala functions. Since both methods are quite similar, I’m going to talk about `andThen` only. Readers who would like to extend to `compose` may try it as a programming exercise.

Method andThen for `Function1[T1, R]` has the following signature:

def andThen[A](g: (R) => A): (T1) => A

A trivial example:

val double: Int => Int = _ * 2
val add1: Int => Int = _ + 1

val doubleThenAdd1 = double andThen add1
// doubleThenAdd1: Int => Int = scala.Function1$Lambda$...

doubleThenAdd1(10)
// res1: Int = 21

Now, let’s replace the 2nd function `add1` with a partial function `inverse`:

val double: Int => Int = _ * 2
val inverse: PartialFunction[Int, Double] = { case x if x != 0 => 1.0 / x }

val doubleThenInverse = double andThen inverse
// doubleThenInverse: Int => Double = scala.Function1$Lambda$...

doubleThenInverse(10)
// res2: Double = 0.05

doubleThenInverse(0)
// scala.MatchError: 0 (of class java.lang.Integer) ...

doubleThenInverse.isDefinedAt(0)
// error: value isDefinedAt is not a member of Int => Double

Note that `doubleThenInverse` still returns a total function even though the composing function is partial. That’s because PartialFunction is a subclass of Function:

trait PartialFunction[-A, +B] extends (A) => B

hence method `andThen` rightfully returns a total function as advertised.

Unfortunately, that’s undesirable as the resulting function lost the `inverse` partial function’s domain information.

Partial function compositions

Method andThen for `PartialFunction[A, C]` has its signature as follows:

def andThen[C](k: (B) => C): PartialFunction[A, C]

Example:

val inverse: PartialFunction[Int, Double] = { case x if x != 0 => 1.0 / x }
val pfMap: PartialFunction[Double, String] = Map(0.1 -> "a",  0.2 -> "b")

val inverseThenPfMap = inverse andThen pfMap
// inverseThenPfMap: PartialFunction[Int,String] = 

inverseThenPfMap(10)
// res1: String = a

inverseThenPfMap(5)
// res2: String = b

inverseThenPfMap.isDefinedAt(10)
// res3: Boolean = true

inverseThenPfMap.isDefinedAt(5)
// res4: Boolean = true

inverseThenPfMap.isDefinedAt(0)
// res5: Boolean = false

// So far so good ... Now, let's try:

inverseThenPfMap(2)
// java.util.NoSuchElementException: key not found: 0.5

inverseThenPfMap.isDefinedAt(2)
// res6: Boolean = false

That works perfectly, since any given element not in the domain of any of the partial functions being composed should have its corresponding element(s) eliminated from the domain of the composed function. In this case, 0.5 is not in the domain of `pfMap`, hence its corresponding element, 2 (which would have been `inverse`-ed to 0.5), should not be in `inverseThenPfMap`’s domain.

Unfortunately, I neglected to mention I’m on Scala 2.13.x. For Scala 2.12 or below, inverseThenPfMap.isDefinedAt(2) would be `true`.

Turning composed functions into a proper partial function

Summarizing what we’ve looked at, there are two issues at hand:

1. If the first function among the functions being composed is a total function, the composed function is a total function, discarding domain information of any subsequent partial functions being composed.
2. Unless you’re on Scala 2.13+, with the first function being a partial function, the resulting composed function is a partial function, but its domain would not embody domain information of any subsequent partial functions being composed.

To tackle the issues, one approach is to leverage implicit conversion by defining a couple of implicit methods that handle composing a partial function on a total function and on a partial function, respectively.

object ComposeFcnOps {
  // Implicit conversion for total function
  implicit class TotalCompose[A, B](f: Function[A, B]) {
    def andThenPF[C](that: PartialFunction[B, C]): PartialFunction[A, C] =
      Function.unlift(x => Option(f(x)).flatMap(that.lift))
  }

  // Implicit conversion for partial function (Not needed on Scala 2.13+)
  implicit class PartialCompose[A, B](pf: PartialFunction[A, B]) {
    def andThenPF[C](that: PartialFunction[B, C]): PartialFunction[A, C] =
      Function.unlift(x => pf.lift(x).flatMap(that.lift))
  }
}

Note that the implicit methods are defined as methods within `implicit class` wrappers, a common practice for the implicit conversion to carry out by invoking the methods like calling class methods.

In the first implicit class, function `f` (i.e. the total function to be implicitly converted) is first transformed to return an `Option`, chained using `flatMap` to the lifted partial function (i.e. the partial function to be composed), followed by an `unlift` to return a partial function.

Similarly, in the second implicit class, function `pf` (i.e. the partial function to be implicitly converted) is first lifted, chained to the lifted partial function (i.e. the partial function to be composed), followed by an `unlift`.

In both cases, `andThenPF` returns a partial function that incorporates the partial domains of the functions involved in the function composition.

Let’s reuse the `double` and `inverse` functions from a previous example:

val double: Int => Int = _ * 2
val inverse: PartialFunction[Int, Double] = { case x if x != 0 => 1.0 / x }

val doubleThenInverse = double andThen inverse
// doubleThenInverse: Int => Double = scala.Function1$Lambda$...

Recall from that example that `doubleThenInverse` is a total function. Now, let’s replace `andThen` with our custom `andThenPF`:

import ComposeFcnOps._

val doubleThenPFInverse: PartialFunction[Int, Double] = double andThenPF inverse
// doubleThenPFInverse: PartialFunction[Int,Double] = 

doubleThenPFInverse(10)
// res1: Double = 0.05

doubleThenPFInverse(0)
// scala.MatchError: 0 (of class java.lang.Integer) ...

doubleThenPFInverse.isDefinedAt(10)
// res2: Boolean = true

doubleThenPFInverse.isDefinedAt(0)
// res2: Boolean = false

The resulting function is now a partial function with the composing function’s partial domain as its own domain. I’ll leave testing for the cases in which the function to be composed is a partial function to the readers.