콘텐츠로 건너뛰기
Home » WWDC의 자동 메모이제이션 코드 분석

WWDC의 자동 메모이제이션 코드 분석

소수 판별이나 소인수분해와 같이 시간이 오래 소요될 수 있는 연산량이 큰 작업이 있고, 이 연산을 반복해서 수행할 일이 많을 때 성능을 개선할 수 있는 가장 손쉬운 방법으로는 연산 결과를 캐싱하는 것이 있을 수 있다. 즉 함수의 외부에 사전과 같은 캐시 저장소를 하나 준비해 두고 입력값을 키로 캐시값을 감사하여 있으면 사용, 없으면 원래 함수 코드를 실행하여 캐시값을 기록하고 리턴하는 방식을 사용하는 것이다.

이는 일일이 그러한 함수를 작성하기 보다는 클로저를 인자로 받아 클로저를 리턴하는 함수를 사용하여 조금 더 유연하게 사용할 수 있다. 이러한 방식을 자동 메모이제이션이라 하는데, 파이썬에서는 흔히 다음과 같은 방법으로 작성한다.

from functools import wraps

def memoization(f):
  m = {}
  @wraps(f)
  def inner(x):
    if x in m:
      return m[x]
    r = f(x)
    m[x] = r
    return r
  return inner

@memoizaiton
def is_prime(n):
  ...

Swift에서도 똑같은 방식으로 자동 메모이제이션을 적용할 수 있다. 대략의 아이디어는 다음과 같다.

  1. (T) -> U 형태의 코드 블럭을 인자로 받는다. 이를 body 라 하자.
  2. Dictionary<T, U> 타입의 사전을 하나 만들어 둔다.
  3. 결과가 될 클로저 result를 정의한다. result의 인자를 x라 하자.
  4. x 가 사전의 키로 존재한다면, 이전에 캐싱된 결과가 있는 것이므로 해당 사전 값을 리턴한다.
  5. 그렇지 않다면 body(x)를 호출하고, 그 결과를 캐싱한 후 리턴한다.
func memoize<T:Hashable, U>(body: @escaping((T)->U)) -> ((T)->U) {
    var memo: Dictionary<T, U> = [:]
    return { (x) in 
        if let res = memo[x] { return res }
        let res = body(x)
        memo[x] = res
        return res
    }
}

let is_prime:(Int)->Bool = { (n) in
  ...
}

이러한 자동 메모이제이션은 대체로 잘 작동하지만, 문제가 하나 있다. 만약 body가 재귀 호출을 하는 경우이다. 실제로 연산을 위해 호출되는 코드는 메모이제이션이 적용된 코드인데, 클로저 내의 body는 메모이제이션이 적용되지 않은 코드이다. 따라서 재귀 호출을 시작한 이후부터는 이미 계산한 적 있는 결과도 불리하게 계산된다.

재귀 함수에도 적용할 수 있는 자동 메모이제이션을 구현하려면 body 내에서 메모이제이션이 적용된 함수를 호출해야 한다. 그런데 “메모이제이션이 적용된 함수”라는 것은 자동 메모이제이션의 결과로, 코드가 작성되는 시점이 아닌 실행되는 시점에 만들어질 함수인데, 이게 가능할까? 일단 body 입장에서는 전달된 파라미터를 사용하는 것이므로 문제가 없을 것이다.

func memoize<T:Hashable, U>(body: @escaping(((T)->U, T)->U)) -> ((T)->U) 
{
    var cache: Dictionary<T, U> = [:]
    var result: ((T)->U)!
    result = { (x) in
        if let res = memo[x] { return res }
        let res = body(result, x)
        memo[x] = res
        return res
    }
    return result
}

let fib:(Int)->Int = { (fib, n) in
    if n < 2 { return 1 }
    return fib(n - 2) + fib(n - 1)    //  여기서 fib 는 파라미터인 fib이고, 메모이제이션된 코드블럭이다.
}

이 코드가 예전에 WWDC에서 소개된 자동 메모이제이션 코드인데, 언뜻 봐서는 어떻게 하자는 것인지 이해가 잘 되지 않고 돌아가는 것이 이상하게 생각되는 것처럼 보인다. 이는 파라미터 이름이 하필 중복되어서 요상해 보이는 것인데, 이렇게 생각하면 조금 더 이해가 쉽다. (근데 코드 써놓고 보니 또 좀 이상해 보이기는 한다….)

let fib:(Int)->Int = { (memoized_fib, n) in 
    if n < 2 { return 1 }
    return memoized_fib(n - 2) + memoized_fib(n - 1)
}