Scala 3 a klíčové slovo inline

2. 6. 2021 #Scala

Nedávno vyšla třetí verze Scaly a přinesla s sebou mnoho velkých změn. Dnes se chci zaměřit jen na jednu z nich: klíčové slovo inline. To umožňuje používat abstrakce s nulovou režií a zároveň představují bránu do světa maker. Scala v minulosti měla asi tři nezávislé verze makrosystému a teprve tenhle vypadá konečně příčetně.

Mě na inline zajímá především garance, že abstrakce budou zcela eliminovány a výsledkem bude stejný bytekód, jako kdyby smyčky byly napsané ručně. Techniky kombinující specializace a type class nejsou na 100% spolehlivé, ale klíčové slovo inline garantuje, že tělo metody bude inlinováno do místa odkud se volá.

S jeho pomocí můžu napsat obecnou inline funkci pro redukci pole

inline def op0[T](init: T, inline f: (T, T) => T, arr: Array[T]): T = {
  var s = init
  var i = 0; while (i < arr.length) {
    s = f(s, arr(i))
    i += 1
  }
  s
}

a pak příslušné konkrétní funkce specializovaná pro daný typ a operaci.

def sum(arr: Array[Double])     = op0(0.0d, _ + _, arr)
def sum(arr: Array[Float])      = op0(0.0f, _ + _, arr)
def product(arr: Array[Double]) = op0(1.0d, _ * _, arr)
def product(arr: Array[Float])  = op0(1.0f, _ * _, arr)

Do každé z nich bude vloženo tělo op0 a výsledek je bytekód bez jakékoli abstrakce. Není v něm žádné volání anonymních funkcí, žádná alokace, nic. Vše inlinováno do podoby prosté smyčky.

inline u argumentu f způsobí, že i tělo funkce předané jako argument bude inlinováno a po volání anonymní funkce nezbudou ani stopy. Díky dvojí inline anotaci v programu nezůstane žádný invokevirtual nebo invokedynamic.

Můžu třeba snadno a bez použití maker vytvořit obdobu klasické for smyčky, která ve Scale jinak chybí. Je u ní garantováno, že poběží stejně rychle jako, kdyby ji jazyk měl a nezáleží, jak úspěšný bude JIT při optimalizaci.

inline def forloop[T](inline init: T)(inline cond: T => Boolean, inline incr: T => T)(inline op: T => Unit) = {
  var i = init; while (cond(i)) {
    op(i)
    i = incr(i)
  }
}

Použití není tak přirozené, jako kdyby šlo o zabudovaný konstrukt, ale lepší než while u níž rád zapomínám na i += 1 a pak čekám, kdy nekonečná smyčka skončí.

def sum(xs: Array[Double]) = {
  var s = 0.0
  // for (var i = 0; i < xs.length; i += 1)
  forloop(0)(_ < xs.length, _ + 1) { i =>
    s += xs(i)
  }
  s
}

Výsledný bytekód je zase prostý jakékoli abstrakce. V tomto případě došlo k eliminování alokace DoubleRef, která je nutná v situacích, kdy closure mění lokální proměnnou s.

Jako o něco komplikovanější příklad vezmu stejnou funkci jako tady. Počítá velikost průniku množin reprezentovaných jako seřazené pole.

def intersectionSize[T: Numeric](a: Array[T], b: Array[T]): Int

Jak ji napsat, aby bylo zaručeno, že vždy poběží stejně efektivně jako v případě ruční specializace pro každý typ?

V odkazovaném článku jsem použil kombinaci specializace a type class. JIT může, pokud je specializovaná nejen funkce sama, ale i metody na type class, vše pěkně inlinovat, ale není garantováno, že vždy uspěje. inline poskytne záruky.

def intersectionSize(a: Array[Short], b: Array[Short]): Int =
  intersectionSizeBody(a, b, ShortOrd)

inline def intersectionSizeBody[T](a: Array[T], b: Array[T], inline ord: Ord[T]): Int = {
  var ai, bi, c = 0;
  while (ai < a.length && bi < b.length) {
    val aa = a(ai)
    val bb = b(bi)
    c  += (if (ord.eq(aa, bb)) 1 else 0);
    ai += (if (ord.lteq(aa, bb)) 1 else 0);
    bi += (if (ord.gteq(aa, bb)) 1 else 0);
  }
  c
}

trait Ord[T] {
  def eq(a: T, b: T): Boolean
  def lteq(a: T, b: T): Boolean
  def gteq(a: T, b: T): Boolean
}

object ShortOrd extends Ord[Short] {
  inline def eq(a: Short, b: Short): Boolean = a == b
  inline def lteq(a: Short, b: Short): Boolean = a <= b
  inline def gteq(a: Short, b: Short): Boolean = a >= b
}

intersectionSize je konkrétní instanciace pro daný typ, inline funkce intersectionSizeBody je expandována všude, kde je použita. Anotace argumentu ord jako inline znamená, že se jeho instance zkopíruje na všechna místa použití. To spolu s inline def metodami na type class Ord způsobí, že jejich těla se plně inlinují do místa volání Ve výsledku je všechna abstrakce kompletně eliminována již na straně scalac.

Nová inline funkce svou silou téměř dosahují síly templatování. Ne úplně, kód musí typechecknout před instanciací.

transparent inline def x0[T](a: T) = a + 1
transparent inline def x1(a: Any) = a + 1

Ve Scale ani jedno neprojde, zatímco klasické templatování àla D by definice povolilo a kompilátor by si stěžoval až v okamžiku instanciace, kdy typ nemá příslušnou metodu.

Templatováním zavání funkce typeChecks kontrolující, jestli je předaný jako string validní v daném (inlinovaném) kontextu.

inline if (scala.compiletime.testing.typeChecks("val x: Short = array(0)")) {}

Tohle (ve spojení s transparent metodami) je překvapivě mocný nástroj, jež bez psaní maker dovoluje staticky podnikat kroky obcházejí typový systém a lidé jistě přijdou na tisíc a jeden způsob jak to využít a zneužít pro fantastické věci.

U příležitosti nové verze jsem taky aktualizoval několik velice starých článků publikovaných dva blogy nazpátek: Zřetězené porovnávání, Klasický for cyklus, Postfixový if, Operátor mocniny, Booleovská kompozice funkcí.

0xDEADBEEF

Scala 3 a klíčové slovo inline