Swarm 0.4 release

The Swarm development team is very proud to announce the latest release of the game. This should still be considered a development/preview release—you still can’t save your games—but it’s made some remarkable progress and there are lots of fun things to try.

What is it?

As a reminder, Swarm is a 2D, open-world programming and resource gathering game with a strongly-typed, functional programming language and a unique upgrade system. Unlocking language features is tied to collecting resources, making it an interesting challenge to bootstrap your way into the use of the full language. It has also become a flexible and powerful platform for constructing programming challenges.

A few of the most significant new features are highlighted below; for full details, see the release notes. If you just want to try it out, see the installation instructions.

Expanded design possibilities

The default play mode is the open-world, resource-gathering scenario—but Swarm also supports “challenge scenarios”, where you have to complete one or more specific objectives with given resources on a custom map. There are currently 58 scenarios and counting—some are silly proofs of concept, but many are quite fun and challenging! I especially recommend checking out the Ranching and Sokoban scenarios, as well as A Frivolous Excursion (pictured below). And creating new scenarios is a great way you can contribute to Swarm even if you don’t know Haskell, or aren’t comfortable hacking on the codebase.

Recently, a large amount of work has gone into expanding the possibilities for scenario design:

  • Structure templates allow you to design map tiles and then reuse them multiple times within a scenario.
  • Waypoints and portals provide a mechanism for automatically navigating and teleporting around the world.
  • Scenarios can have multiple subworlds besides the main “overworld”, connected by portals. For example you could go “into a building” and have a separate map for the building interior.
  • There are a slew of new robot commands, many to do with different sensing modalities: stride, detect, sniff, chirp, resonate, watch, surveil, scout, instant, push, density, use, halt, and backup.
  • A new domain-specific language for describing procedurally generated worlds. The default procedurally generated world used to be hardcoded, but now it is described externally via the new DSL, and you can design your own procedurally generated worlds without editing the Swarm source code.
  • The key input handler feature allows you to program robots to respond to keyboard input, so you can e.g. drive them around manually, or interactively trigger more complex behaviors. This makes it possible to design “arcade-style” challenges, where the player needs to guide a robot and react to obstacles in real time—but they get to program the robot to respond to their commands first!
  • A new prototype integrated world editor lets you design worlds interactively.

UI improvements

In the past, entity and goal descriptions were simply plain text; recently, we switched to actually parsing Markdown. Partly, this is just to make things look nice, since we can highlight code snippets, entity names, etc.:

But it also means that we can now validate all code examples and entity names, and even test that the tutorial is pedagogically sound: any command used in a tutorial solution must be mentioned in a previous tutorial, or else our CI fails!

There are also a number of other small UI enhancements, such as improved type error messages, inventory search, and a collapsible REPL panel, among others.

Scoring metrics

We now keep track of a number of metrics related to challenge scenario solutions, such as total time, total game ticks, and code size. These metrics are tracked and saved across runs, so you can compete with yourself, and with others. For now, see these wiki pages:

In the future, perhaps there will eventually be some kind of social website with leaderboards and user-uploaded scenarios.

Debugging

Last but not least, we now have an integrated single-stepping and debugging mode (enabled by the tweezers device).

Give it a try!

To install, check out the installation instructions: you can download a binary release (for now, Linux only, but MacOS binaries should be on the horizon), or install from Hackage. Give it a try and send us your feedback, either via a github issue or IRC!

Future plans & getting involved

We’re still hard at work on the game. Fun upcoming things include:

Of course, there are also tons of small things that need fixing and polishing too! If you’re interested in getting involved, check out our contribution guide, come join us on IRC (#swarm on Libera.Chat), or take a look at the list of issues marked “low-hanging fruit”.

Brought to you by the Swarm development team:

  • Brent Yorgey
  • Karl Ostmo
  • Ondřej Šebek

With contributions from:

  • Alexander Block
  • Brian Wignall
  • Chris Casinghino
  • Daniel Díaz Carrete
  • Huw Campbell
  • Ishan Bhanuka
  • Jacob
  • Jens Petersen
  • José Rafael Vieira
  • Joshua Price
  • lsmor
  • Noah Yorgey
  • Norbert Dzikowski
  • Paul Brauner
  • Ryan Yates
  • Sam Tay
  • Steven Garcia
  • Tamas Zsar
  • Tristan de Cacqueray
  • Valentin Golev

…not to mention many others who gave valuable suggestions and feedback. Want to see your name listed here in the next release? See how you can contribute!

Posted in haskell, projects | Tagged , , , , | Leave a comment

Compiling to Intrinsically Typed Combinators

tl;dr: How to compile a functional language via combinators (and evaluate via the Haskell runtime) while keeping the entire process type-indexed, with a bibliography and lots of references for further reading

There is a long history, starting with Schönfinkel and Curry, of abstracting away variable names from lambda calculus terms by converting to combinators, aka bracket abstraction. This was popular in the 80’s as a compilation technique for functional languages (Turner, 1979; Augustsson, 1986; Jones, 1987; Diller, 1988), then apparently abandoned. More recently, however, it has been making a bit of a comeback. For example, see Naylor (2008), Gratzer (2015), Lynn (2017), and Mahler (2021). Bracket abstraction is intimately related to compiling to cartesian closed categories (Elliott, 2017; Mahler, 2021), and also enables cool tricks like doing evaluation via the Haskell runtime system (Naylor, 2008; Seo, 2016; Mahler, 2022).

However, it always bothered me that the conversion to combinators was invariably described in an untyped way. Partly to gain some assurance that we are doing things correctly, but mostly for fun, I wondered if it would be possible to do the whole pipeline in an explicitly type-indexed way. I eventually found a nice paper by Oleg Kiselyov (2018) which explains exactly how to do it (it even came with OCaml code that I was easily able to port to Haskell!).

In this blog post, I:

  • Show an example of typechecking and elaboration for a functional language into explicitly type-indexed terms, such that it is impossible to write down ill-typed terms
  • Demonstrate a Haskell port of Oleg Kiselyov’s typed bracket abstraction algorithm
  • Demonstrate type-indexed evaluation of terms via the Haskell runtime
  • Put together an extensive bibliography with references for further reading

This blog post is rendered automatically from a literate Haskell file; you can find the complete working source code and blog post on GitHub. I’m always happy to receive comments, fixes, or suggestions for improvement.

But First, A Message From Our Sponsors

So many yummy language extensions.

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UnicodeSyntax #-}
{-# LANGUAGE ViewPatterns #-}

module TypedCombinators where

import Control.Monad.Combinators.Expr
import Data.Functor.Const qualified as F
import Data.Void
import Data.Text ( Text )
import Data.Text qualified as T
import Data.Kind (Type)
import Data.Type.Equality ( type (:~:)(Refl), TestEquality(..) )
import Text.Megaparsec
import Text.Megaparsec.Char
import Text.Megaparsec.Char.Lexer qualified as L
import Witch (into)
import Prelude hiding (lookup)

Raw terms and types

Here’s an algebraic data type to represent raw terms of our DSL, something which might come directly out of a parser. The exact language we use here isn’t all that important; I’ve put in just enough features to make it nontrivial, but not much beyond that. We have integer literals, variables, lambdas, application, let and if expressions, addition, and comparison with >. Of course, it would be easy to add more types, constants, and language features.

data Term where
  Lit :: Int -> Term
  Var :: Text -> Term
  Lam :: Text -> Ty -> Term -> Term
  App :: Term -> Term -> Term
  Let :: Text -> Term -> Term -> Term
  If  :: Term -> Term -> Term -> Term
  Add :: Term -> Term -> Term
  Gt  :: Term -> Term -> Term
  deriving Show

A few things to note:

  • In order to keep things simple, notice that lambdas must be annotated with the type of the argument. There are other choices we could make, but this is the simplest for now. I’ll have more to say about other choices later.

  • I included if not only because it gives us something to do with Booleans, but also because it is polymorphic, which adds an interesting twist to our typechecking.

  • I included >, not only because it gives us a way to produce Boolean values, but also because it uses ad-hoc polymorphism, that is, we can compare at any type which is an instance of Ord. This is an even more interesting twist.

Here are our types: integers, booleans, and functions.

data Ty where
  TyInt  :: Ty
  TyBool :: Ty
  TyFun  :: Ty -> Ty -> Ty
  deriving Show

Finally, here’s an example term that uses all the features of our language (I’ve included a simple parser in an appendix at the end of this post):

example :: Term
example = readTerm $ T.unlines
  [ "let twice = \\f:Int -> Int. \\x:Int. f (f x) in"
  , "let z = 1 in"
  , "if 7 > twice (\\x:Int. x + 3) z then z else z + 1"
  ]

Since 7 is not, in fact, strictly greater than 1 + 3 + 3, this should evaluate to 2.

Type-indexed constants

That was the end of our raw, untyped representations—from now on, everything is going to be type-indexed! First of all, we’ll declare an enumeration of constants, with each constant indexed by its corresponding host language type. These will include both any special language built-ins (like if, +, and >) as well as a set of combinators which we’ll be using as a compilation target—more on these later.

data Const :: Type -> Type where
  CInt :: Int -> Const Int
  CIf :: Const (Bool -> α -> α -> α)
  CAdd :: Const (Int -> Int -> Int)
  CGt :: Ord α => Const (α -> α -> Bool)
  I :: Const (α -> α)
  K :: Const (α -> b -> α)
  S :: Const ((α -> b -> c) -> (α -> b) -> α -> c)
  B :: Const ((     b -> c) -> (α -> b) -> α -> c)
  C :: Const ((α -> b -> c) ->       b  -> α -> c)

deriving instance Show (Const α)

The polymorphism of if (and the combinators I, K, etc., for that matter) poses no real problems. If we really wanted the type of CIf to be indexed by the exact type of if, it would be something like

  CIf :: Const ( α. Bool -> α -> α -> α)

but this would require impredicative types which can be something of a minefield. However, what we actually get is

  CIf ::  α. Const (Bool -> α -> α -> α)

which is unproblematic and works just as well for our purposes.

The type of CGt is more interesting: it includes an Ord α constraint. That means that at the time we construct a CGt value, we must have in scope an Ord instance for whatever type α is; conversely, when we pattern-match on CGt, we will bring that instance into scope. We will see how to deal with this later.

For convenience, we make a type class HasConst for type-indexed things that can contain embedded constants (we will end up with several instances of this class).

class HasConst t where
  embed :: Const α -> t α

Also for convenience, here’s a type class for type-indexed things that support some kind of application operation. (Note that we don’t necessarily want to require t to support a pure :: a -> t a operation, or even be a Functor, so using Applicative would not be appropriate, even though $$ has the same type as <*>.)

infixl 1 $$
class Applicable t where
  ($$) :: t (α -> β) -> t α -> t β

Note that, unlike the standard $ operator, $$ is left-associative, so, for example, f $$ x $$ y should be read just like f x y, that is, f $$ x $$ y = (f $$ x) $$ y.

Finally, we’ll spend a bunch of time applying constants to things, or applying things to constants, so here are a few convenience operators for combining $$ and embed:

infixl 1 .$$
(.$$) :: (HasConst t, Applicable t) => Const (α -> β) -> t α -> t β
c .$$ t = embed c $$ t

infixl 1 $$.
($$.) :: (HasConst t, Applicable t) => t (α -> β) -> Const α -> t β
t $$. c = t $$ embed c

infixl 1 .$$.
(.$$.) :: (HasConst t, Applicable t) => Const (α -> β) -> Const α -> t β
c1 .$$. c2 = embed c1 $$ embed c2

Type-indexed types and terms

Now let’s build up our type-indexed core language. First, we’ll need a data type for type-indexed de Bruijn indices. A value of type Idx γ α is a variable with type α in the context γ (represented as a type-level list of types). For example, Idx [Int,Bool,Int] Int would represent a variable of type Int (and hence must either be variable 0 or 2).

data Idx :: [Type] -> Type -> Type where
  VZ :: Idx (α ': γ) α
  VS :: Idx γ α -> Idx (β ': γ) α

deriving instance Show (Idx γ α)

Now we can build our type-indexed terms. Just like variables, terms are indexed by a typing context and a type; t : TTerm γ α can be read as “t is a term with type α, possibly containing variables whose types are described by the context γ”. Our core language has only variables, constants, lambdas, and application. Note we’re not just making a type-indexed version of our original term language; for simplicity, we’re going to simultaneously typecheck and elaborate down to this much simpler core language. (Of course, it would also be entirely possible to introduce another intermediate data type for type-indexed terms, and separate the typechecking and elaboration phases.)

data TTerm :: [Type] -> Type -> Type where
  TVar :: Idx γ α -> TTerm γ α
  TConst :: Const α -> TTerm γ α
  TLam :: TTerm (α ': γ) β -> TTerm γ (α -> β)
  TApp :: TTerm γ (α -> β) -> TTerm γ α -> TTerm γ β

deriving instance Show (TTerm γ α)

instance Applicable (TTerm γ) where
  ($$) = TApp

instance HasConst (TTerm γ) where
  embed = TConst

Now for some type-indexed types!

data TTy :: Type -> Type where
  TTyInt :: TTy Int
  TTyBool :: TTy Bool
  (:->:) :: TTy α -> TTy β -> TTy (α -> β)

deriving instance Show (TTy ty)

TTy is a term-level representation of our DSL’s types, indexed by corresponding host language types. In other words, TTy is a singleton: for a given type α there is a single value of type TTy α. Put another way, pattern-matching on a value of type TTy α lets us learn what the type α is. (See (Le, 2017) for a nice introduction to the idea of singleton types.)

We will need to be able to test two value-level type representations for equality and have that reflected at the level of type indices; the TestEquality class from Data.Type.Equality is perfect for this. The testEquality function takes two type-indexed things and returns a type equality proof wrapped in Maybe.

instance TestEquality TTy where
  testEquality :: TTy α -> TTy β -> Maybe (α :~: β)
  testEquality TTyInt TTyInt = Just Refl
  testEquality TTyBool TTyBool = Just Refl
  testEquality (α₁ :->: β₁) (α₂ :->: β₂) =
    case (testEquality α₁ α₂, testEquality β₁ β₂) of
      (Just Refl, Just Refl) -> Just Refl
      _ -> Nothing
  testEquality _ _ = Nothing

Recall that the CGt constant requires an Ord instance; the checkOrd function pattern-matches on a TTy and witnesses the fact that the corresponding host-language type has an Ord instance (if, in fact, it does).

checkOrd :: TTy α -> (Ord α => r) -> Maybe r
checkOrd TTyInt r = Just r
checkOrd TTyBool r = Just r
checkOrd _ _ = Nothing

As a quick aside, for simplicity’s sake, I am going to use Maybe throughout the rest of this post to indicate possible failure. In a real implementation, one would of course want to return more information about any error(s) that occur.

Existential wrappers

Sometimes we will need to wrap type-indexed things inside an existential wrapper to hide the type index. For example, when converting from a Ty to a TTy, or when running type inference, we can’t know in advance which type we’re going to get. So we create the Some data type which wraps up a type-indexed thing along with a corresponding TTy. Pattern-matching on the singleton TTy will allow us to recover the type information later.

data Some :: (Type -> Type) -> Type where
  Some :: TTy α -> t α -> Some t

mapSome :: ( α. s α -> t α) -> Some s -> Some t
mapSome f (Some α t) = Some α (f t)

The first instantiation we’ll create is an existentially wrapped type, where the TTy itself is the only thing we care about, and the corresponding t will just be the constant unit type functor. It would be annoying to keep writing F.Const () everywhere so we create some type and pattern synonyms for convenience.

type SomeTy = Some (F.Const ())

pattern SomeTy :: TTy α -> SomeTy
pattern SomeTy α = Some α (F.Const ())
{-# COMPLETE SomeTy #-}

The someType function converts from a raw Ty to a type-indexed TTy, wrapped up in an existential wrapper.

someType :: Ty -> SomeTy
someType TyInt = SomeTy TTyInt
someType TyBool = SomeTy TTyBool
someType (TyFun a b) = case (someType a, someType b) of
  (SomeTy α, SomeTy β) -> SomeTy (α :->: β)

Type inference and elaboration

Now that we have our type-indexed core language all set, it’s time to do type inference, that is, translate from untyped terms to type-indexed ones! First, let’s define type contexts, i.e. mappings from variables to their types. We store contexts simply as a (fancy, type-indexed) list of variable names paired with their types. This is inefficient—it takes linear time to do a lookup—but we don’t care, because this is an intermediate representation used only during typechecking. By the time we actually get around to running terms, variables won’t even exist any more.

data Ctx :: [Type] -> Type where

  -- CNil represents an empty context.
  CNil :: Ctx '[]

  -- A cons stores a variable name and its type,
  -- and then the rest of the context.
  (:::) :: (Text, TTy α) -> Ctx γ -> Ctx (α ': γ)

Now we can define the lookup function, which takes a variable name and a context and tries to return a corresponding de Bruijn index into the context. When looking up a variable name in the context, we can’t know in advance what index we will get and what type it will have, so we wrap the returned Idx in Some.

lookup :: Text -> Ctx γ -> Maybe (Some (Idx γ))
lookup _ CNil = Nothing
lookup x ((y, α) ::: ctx)
  | x == y = Just (Some α VZ)
  | otherwise = mapSome VS <$> lookup x ctx

Now we’re finally ready to define the infer function! It takes a type context and a raw term, and tries to compute a corresponding type-indexed term. Note that there’s no particular guarantee that the term we return corresponds to the input term—we will just have to be careful—but at least the Haskell type system guarantees that we can’t return a type-incorrect term, which is especially important when we have some nontrivial elaboration to do. Of course, just as with variable lookups, when inferring the type of a term we can’t know in advance what type it will have, so we will need to return an existential wrapper around a type-indexed term.

infer :: Ctx γ -> Term -> Maybe (Some (TTerm γ))
infer ctx = \case

To infer the type of a literal integer value, just return TTyInt with a literal integer constant.

  Lit i -> return $ Some TTyInt (embed (CInt i))

To infer the type of a variable, look it up in the context and wrap the result in TVar. Notice how we are allowed to pattern-match on the Some returned from lookup (revealing the existentially quantified type inside) since we immediately wrap it back up in another Some when returning the TVar.

  Var x -> mapSome TVar <$> lookup x ctx

To infer the type of a lambda, we convert the argument type annotation to a type-indexed type, infer the type of the body under an extended context, and then return a lambda with an appropriate function type. (If lambdas weren’t required to have type annotations, then we would either have to move the lambda case to the check function, or else use unification variables and solve type equality constraints. The former would be straightforward, but I don’t know how to do the latter in a type-indexed way—sounds like a fun problem for later.)

  Lam x a t -> do
    case someType a of
      Some α _ -> do
        Some β t' <- infer ((x,α) ::: ctx) t
        return $ Some (α :->: β) (TLam t')

To infer the type of an application, we infer the type of the left-hand side, ensure it is a function type, and check that the right-hand side has the correct type. We will see the check function later.

  App t1 t2 -> do
    Some τ t1' <- infer ctx t1
    case τ of
      α :->: β -> do
        t2' <- check ctx α t2
        return $ Some β (TApp t1' t2')
      _ -> Nothing

To infer the type of a let-expression, we infer the type of the definition, infer the type of the body under an extended context, and then desugar it into an application of a lambda. That is, let x = t1 in t2 desugars to (\x.t2) t1.

  Let x t1 t2 -> do
    Some α t1' <- infer ctx t1
    Some β t2' <- infer ((x, α) ::: ctx) t2
    return $ Some β (TApp (TLam t2') t1')

Note again that we can’t accidentally get mixed up here—for example, if we incorrectly desugar to (\x.t1) t2 we get a Haskell type error, like this:

    • Couldn't match type ‘γ’ with ‘α : γ’
      Expected: TTerm γ α1
        Actual: TTerm (α : γ) α1

To infer an if-expression, we can check that the test has type Bool, infer the types of the two branches, and ensure that they are the same. If so, we return the CIf constant applied to the three arguments. The reason this typechecks is that pattern-matching on the Refl from the testEquality call brings into scope the fact that the types of t2 and t3 are equal, so we can apply CIf which requires them to be so.

  If t1 t2 t3 -> do
    t1' <- check ctx TTyBool t1
    Some α t2' <- infer ctx t2
    Some β t3' <- infer ctx t3
    case testEquality α β of
      Nothing -> Nothing
      Just Refl -> return $ Some α (CIf .$$ t1' $$ t2' $$ t3')

Addition is simple; we just check that both arguments have type Int.

  Add t1 t2 -> do
    t1' <- check ctx TTyInt t1
    t2' <- check ctx TTyInt t2
    return $ Some TTyInt (CAdd .$$ t1' $$ t2')

“Greater than” is a bit interesting because we allow it to be used at both Int and Bool. So, just as with if, we must infer the types of the arguments and check that they match. But then we must also use the checkOrd function to ensure that the argument types are an instance of Ord. In particular, we wrap CGt (which requires an Ord constraint) in a call to checkOrd α (which provides one).

  Gt t1 t2 -> do
    Some α t1' <- infer ctx t1
    Some β t2' <- infer ctx t2
    case testEquality α β of
      Nothing -> Nothing
      Just Refl -> (\c -> Some TTyBool (c .$$ t1' $$ t2')) <$> checkOrd α CGt

Finally, here’s the check function: to check that an expression has an expected type, just infer its type and make sure it’s the one we expected. (With more interesting languages we might also have more cases here for terms which can be checked but not inferred.) Notice how this also allows us to return the type-indexed term without using an existential wrapper, since the expected type is an input.

check :: Ctx γ -> TTy α -> Term -> Maybe (TTerm γ α)
check ctx α t = do
  Some β t' <- infer ctx t
  case testEquality α β of
    Nothing -> Nothing
    Just Refl -> Just t'

Putting this all together so far, we can check that the example term has type Int and see what it elaborates to (I’ve included a simple pretty-printer for TTerm in an appendix):

λ> putStrLn . pretty . fromJust . check CNil TTyInt $ example
(λ. (λ. if (gt 7 (x1 (λ. plus x0 3) x0)) x0 (plus x0 1)) 1) (λ. λ. x1 (x1 x0))

An aside: a typed interpreter

We can now easily write an interpreter. However, this is pretty inefficient (it has to carry around an environment and do linear-time variable lookups), and later we’re going to compile our terms directly to host language terms. So this interpreter is just a nice aside, for fun and testing.

With that said, given a closed term, we can interpret it directly to a value of its corresponding host language type. We need typed environments and a indexing function (note that for some reason GHC can’t see that the last case of the indexing function is impossible; if we tried implementing it in, say, Agda, we wouldn’t have to write that case).

data Env :: [Type] -> Type where
  ENil :: Env '[]
  ECons :: α -> Env γ -> Env (α ': γ)

(!) :: Env γ -> Idx γ α -> α
(ECons x _) ! VZ = x
(ECons _ e) ! (VS x) = e ! x
ENil ! _ = error "GHC can't tell this is impossible"

Now the interpreter is straightforward. Look how beautifully everything works out with the type indexing.

interpTTerm :: TTerm '[] α -> α
interpTTerm = go ENil
  where
    go :: Env γ -> TTerm γ α -> α
    go e = \case
      TVar x -> e ! x
      TLam body -> \x -> go (ECons x e) body
      TApp f x -> go e f (go e x)
      TConst c -> interpConst c

interpConst :: Const α -> α
interpConst = \case
  CInt i -> i
  CIf -> \b t e -> if b then t else e
  CAdd -> (+)
  CGt -> (>)
  K -> const
  S -> (<*>)
  I -> id
  B -> (.)
  C -> flip
λ> interpTTerm . fromJust . check CNil TTyInt $ example
2

Compiling to combinators: type-indexed bracket abstraction

Now, on with the main attraction! It’s well-known that certain sets of combinators are Turing-complete: for example, SKI is the most well-known complete set (or just SK if you’re trying to be minimal). There are well-known algorithms for compiling lambda calculus terms into combinators, known generally as bracket abstraction (for further reading about bracket abstraction in general, see Diller (2014); for some in-depth history along with illustrative Haskell code, see Ben Lynn’s page on Combinatory Logic (2022); for nice example implementations in Haskell, see blog posts by Gratzer (2015), Seo (2016), and Mahler (2021).)

So the idea is to compile our typed core language down to combinators. The resulting terms will have no lambdas or variables—only constants and application! The point is that by making environments implicit, with a few more tricks we can make use of the host language runtime’s ability to do beta reduction, which will be much more efficient than our interpreter.

The BTerm type below will be the compilation target. Again for illustration and/or debugging we can easily write a direct interpreter for BTerm—but this still isn’t the intended code path. There will still be one more step to convert BTerms directly into host language terms.

data BTerm :: Type -> Type where
  BApp :: BTerm (α -> β) -> BTerm α -> BTerm β
  BConst :: Const α -> BTerm α

deriving instance Show (BTerm ty)

instance Applicable BTerm where
  ($$) = BApp

instance HasConst BTerm where
  embed = BConst

interpBTerm :: BTerm ty -> ty
interpBTerm (BApp f x) = interpBTerm f (interpBTerm x)
interpBTerm (BConst c) = interpConst c

We will use the usual SKI combinators as well as B and C, which are like special-case variants of S:

  • S x y z = x z (y z)
  • B x y z = x (y z)
  • C x y z = x z (y )

S handles the application of x to y in the case where they both need access to a shared parameter z; B and C are similar, but B is used when only y, and not x, needs access to z, and C is for when only x needs access to z. Using B and C will allow for more efficient encodings than would be possible with S alone. If you want to compile a language with recursion you can also easily add the usual Y combinator (“SICKBY”), although the example language in this post has no recursion so we won’t use it.

Bracket abstraction is often presented in an untyped way, but I found this really cool paper by Oleg Kiselyov (2018) where he shows how to do bracket abstraction in a completely compositional, type-indexed way. I found the paper a bit hard to understand, but fortunately it came with working OCaml code! Translating it to Haskell was straightforward. Much later, after writing most of this blog post, I found a a nice explanation of Kiselyov’s paper by Lynn (2022) which helped me make more sense of the paper.

First, a data type for open terms, which represent an intermediate stage in the bracket abstraction algorithm, where some parts have been converted to closed combinator terms (the E constructor embeds BTerm values), and some parts still have not. This corresponds to Kiselyov’s eta-optimized version (section 4.1 of the paper). A simplified version that does not include V is possible, but results in longer combinator expressions.

data OTerm :: [Type] -> Type -> Type where

  -- E contains embedded closed (i.e. already abstracted) terms.
  E :: BTerm α -> OTerm γ α

  -- V represents a reference to the innermost/top environment
  -- variable, i.e. Z
  V :: OTerm (α ': γ) α

  -- N represents internalizing the innermost bound variable as a
  -- function argument. In other words, we can represent an open
  -- term referring to a certain variable as a function which
  -- takes that variable as an argument.
  N :: OTerm γ (α -> β) -> OTerm (α ': γ) β

  -- For efficiency, there is also a special variant of N for the
  -- case where the term does not refer to the topmost variable at
  -- all.
  W :: OTerm γ β -> OTerm (α ': γ) β

instance HasConst (OTerm γ) where
  embed = E . embed

Now for the bracket abstraction algorithm. First, a function to do type- and environment-preserving conversion from TTerm to OTerm. The conv function handles the variable, lambda, and constant cases. The application case is handled by the Applicable instance.

conv :: TTerm γ α -> OTerm γ α
conv = \case
  TVar VZ -> V
  TVar (VS x) -> W (conv (TVar x))
  TLam t -> case conv t of
    V -> E (embed I)
    E d -> E (K .$$ d)
    N e -> e
    W e -> K .$$ e
  TApp t1 t2 -> conv t1 $$ conv t2
  TConst c -> embed c

The Applicable instance for OTerm has 15 cases—one for each combination of OTerm constructors. Why not 16, you ask? Because the V $$ V case is impossible (exercise for the reader: why?). The cool thing is that GHC can tell that case would be ill-typed, and agrees that this definition is total—that is, it does not give a non-exhaustive pattern match warning. This is a lot of code, but understanding each individual case is not too hard if you understand the meaning of the constructors E, V, N, and W. For example, if we have one term that ignores the innermost bound variable being applied to another term that also ignores the innermost bound variable (W e1 $$ W e2), we can apply one term to the other and wrap the result in W again (W (e1 $$ e2)). Other cases use the combinators B, C, S to route the input to the proper places in an application.

instance Applicable (OTerm γ) where
  ($$) :: OTerm γ (α -> β) -> OTerm γ α -> OTerm γ β
  W e1 $$ W e2 = W (e1 $$ e2)
  W e $$ E d = W (e $$ E d)
  E d $$ W e = W (E d $$ e)
  W e $$ V = N e
  V $$ W e = N (E (C .$$. I) $$ e)
  W e1 $$ N e2 = N (B .$$ e1 $$ e2)
  N e1 $$ W e2 = N (C .$$ e1 $$ e2)
  N e1 $$ N e2 = N (S .$$ e1 $$ e2)
  N e $$ V = N (S .$$ e $$. I)
  V $$ N e = N (E (S .$$. I) $$ e)
  E d $$ N e = N (E (B .$$ d) $$ e)
  E d $$ V = N (E d)
  V $$ E d = N (E (C .$$. I $$ d))
  N e $$ E d = N (E (C .$$. C $$ d) $$ e)
  E d1 $$ E d2 = E (d1 $$ d2)

The final bracket abstraction algorithm consists of calling conv on a closed TTerm—this must result in a term of type OTerm '[] α, and the only constructor which could possibly produce such a type is E, containing an embedded BTerm. So we can just extract that BTerm, and GHC can see that this is total.

bracket :: TTerm '[] α -> BTerm α
bracket t = case conv t of { E t' -> t' }

Let’s apply this to our example term and see what we get:

λ> putStrLn . pretty . bracket . fromJust . check CNil TTyInt $ example
C C 1 (C C (C C 1 plus) (B S (C C I (B S (B (B if) (B (B (gt 7)) (C I (C C 3 plus)))))))) (S B I)
λ> interpBTerm . bracket . fromJust . check CNil TTyInt $ example
2

Neat! This is not too much longer than the original term, which is the point of using the optimized version. Interestingly, this example happens to not use K at all, but a more complex term certainly would.

Kiselyov also presents an even better algorithm using n-ary combinators which uses guaranteed linear time and space. For simplicity, he presents it in an untyped way and claims in passing that it “can be backported to the typed case”, though I am not aware of anyone who has actually done this yet (perhaps I will, later). Lynn (2022) has a nice explanation of Kiselyov’s paper, including a section that explores several alternatives to Kiselyov’s linear-time algorithm.

Compiling type-indexed combinators to Haskell

So at this point we can take a Term, typecheck it to produce a TTerm, then use bracket abstraction to convert that to a BTerm. We have an interpreter for BTerms, but we’re instead going to do one more compilation step, to turn BTerms directly into native Haskell values. This idea originates with Naylor (2008) and is well-explained in blog posts by Seo (2016) and Mahler (2022). This still feels a little like black magic to me, and I am actually unclear on whether it is really faster than calling interpBTerm; some benchmarking would be needed. In any case I include it here for completeness.

Our target for this final compilation step is the following CTerm type, which has only functions, represented by CFun, and constants. Note, however, that CConst is intended to be used only for non-function types, i.e. base types, although there’s no nice way (that I know of, at least) to use the Haskell type system to enforce this.

data CTerm α where
  CFun :: (CTerm α -> CTerm β) -> CTerm (α -> β)
  CConst :: α -> CTerm α -- CConst invariant: α is not a function type

instance Applicable CTerm where
  CFun f $$ x = f x
  CConst _ $$ _ = error "CConst should never contain a function!"

compile :: BTerm α -> CTerm α
compile (BApp b1 b2) = compile b1 $$ compile b2
compile (BConst c) = compileConst c

compileConst :: Const α -> CTerm α
compileConst = \case
  (CInt i) -> CConst i
  CIf      -> CFun $ \(CConst b) -> CFun $ \t -> CFun $ \e -> if b then t else e
  CAdd     -> binary (+)
  CGt      -> binary (>)
  K        -> CFun $ \x -> CFun $ \_ -> x
  S        -> CFun $ \f -> CFun $ \g -> CFun $ \x -> f $$ x $$ (g $$ x)
  I        -> CFun id
  B        -> CFun $ \f -> CFun $ \g -> CFun $ \x -> f $$ (g $$ x)
  C        -> CFun $ \f -> CFun $ \x -> CFun $ \y -> f $$ y $$ x

binary :: (α -> b -> c) -> CTerm (α -> b -> c)
binary op = CFun $ \(CConst x) -> CFun $ \(CConst y) -> CConst (op x y)

Finally, we can “run” a CTerm α to extract a value of type α. Typically, if α is some kind of base type like Int, runCTerm doesn’t actually do any work—all the work is done by the Haskell runtime itself. However, for completeness, I include a case for CFun as well.

runCTerm :: CTerm α -> α
runCTerm (CConst a) = a
runCTerm (CFun f) = runCTerm . f . CConst

We can put this all together into our final pipeline:

evalInt :: Term -> Maybe Int
evalInt = fmap (runCTerm . compile . bracket) . check CNil TTyInt
λ> evalInt example
Just 2

Appendices

There’s nothing interesting to see here—unless you’ve never written a parser or pretty-printer before, in which case perhaps it is very interesting! If you want to learn how to write parsers, see this very nice Megaparsec tutorial. And see here for some help writing a basic pretty-printer.

Parsing

type Parser = Parsec Void Text
type ParserError = ParseErrorBundle Text Void

reservedWords :: [Text]
reservedWords = ["let", "in", "if", "then", "else", "Int", "Bool"]

sc :: Parser ()
sc = L.space space1 (L.skipLineComment "--") empty

lexeme :: Parser a -> Parser a
lexeme = L.lexeme sc

symbol :: Text -> Parser Text
symbol = L.symbol sc

reserved :: Text -> Parser ()
reserved w = (lexeme . try) $ string' w *> notFollowedBy alphaNumChar

identifier :: Parser Text
identifier = (lexeme . try) (p >>= nonReserved) <?> "variable name"
 where
  p = (:) <$> letterChar <*> many alphaNumChar
  nonReserved (into @Text -> t)
    | t `elem` reservedWords =
        fail . into @String $
          T.concat ["reserved word '", t, "' cannot be used as variable name"]
    | otherwise = return t

integer :: Parser Int
integer = lexeme L.decimal

parens :: Parser a -> Parser a
parens = between (symbol "(") (symbol ")")

parseTermAtom :: Parser Term
parseTermAtom =
      Lit <$> integer
  <|> Var <$> identifier
  <|> Lam <$> (symbol "\\" *> identifier)
          <*> (symbol ":" *> parseType)
          <*> (symbol "." *> parseTerm)
  <|> Let <$> (reserved "let" *> identifier)
          <*> (symbol "=" *> parseTerm)
          <*> (reserved "in" *> parseTerm)
  <|> If  <$> (reserved "if" *> parseTerm)
          <*> (reserved "then" *> parseTerm)
          <*> (reserved "else" *> parseTerm)
  <|> parens parseTerm

parseTerm :: Parser Term
parseTerm = makeExprParser parseTermAtom
  [ [InfixL (App <$ symbol "")]
  , [InfixL (Add <$ symbol "+")]
  , [InfixL (Gt <$ symbol ">")]
  ]

parseTypeAtom :: Parser Ty
parseTypeAtom =
  TyInt <$ reserved "Int"
  <|> TyBool <$ reserved "Bool"
  <|> parens parseType

parseType :: Parser Ty
parseType = makeExprParser parseTypeAtom
  [ [InfixR (TyFun <$ symbol "->")] ]

readTerm :: Text -> Term
readTerm = either undefined id . runParser parseTerm ""

Pretty-printing

type Prec = Int

class Pretty p where
  pretty :: p -> String
  pretty = prettyPrec 0

  prettyPrec :: Prec -> p -> String
  prettyPrec _ = pretty

mparens :: Bool -> String -> String
mparens True  = ("("++) . (++")")
mparens False = id

instance Pretty (Const α) where
  prettyPrec _ = \case
    CInt i -> show i
    CIf -> "if"
    CAdd -> "plus"
    CGt  -> "gt"
    c -> show c

instance Pretty (Idx γ α) where
  prettyPrec _ = ("x" ++) . show . toNat
    where
      toNat :: Idx γ α -> Int
      toNat VZ = 0
      toNat (VS i) = 1 + toNat i

instance Pretty (TTerm γ α) where
  prettyPrec p = \case
    TVar x -> pretty x
    TConst c -> pretty c
    TLam t -> mparens (p>0) $ "λ. " ++ prettyPrec 0 t
    TApp t1 t2 -> mparens (p>1) $
      prettyPrec 1 t1 ++ " " ++ prettyPrec 2 t2

instance Pretty (BTerm α) where
  prettyPrec p = \case
    BConst c -> pretty c
    BApp t1 t2 -> mparens (p>0) $
      prettyPrec 0 t1 ++ " " ++ prettyPrec 1 t2

References

Augustsson, L. (1986) SMALL: A small interactive functional system. Chalmers Tekniska Högskola/Göteborgs Universitet. Programming Methodology Group.
Diller, A. (1988) Compiling functional languages. John Wiley & Sons, Inc.
Diller, A. (2014) Bracket abstraction algorithms. Available at: https://www.cantab.net/users/antoni.diller/brackets/intro.html.
Elliott, C. (2017) ‘Compiling to categories’, Proceedings of the ACM on Programming Languages, 1(ICFP), pp. 1–27.
Gratzer, D. (2015) Bracket Abstraction: The Smallest PL You’ve Ever Seen. Available at: https://jozefg.bitbucket.io/posts/2015-05-01-brackets.html.
Jones, S.L.P. (1987) The Implementation of Functional Programming Languages. Prentice-Hall, Inc.
Kiselyov, O. (2018) λ to ski, semantically: Declarative pearl’, in International symposium on functional and logic programming. Springer, pp. 33–50.
Le, J. (2017) Introduction to Singletons (Part 1). Available at: https://blog.jle.im/entry/introduction-to-singletons-1.html.
Lynn, B. (2017) A Combinatory Compiler. Available at: https://crypto.stanford.edu/~blynn/lambda/sk.html.
Lynn, B. (2022) Combinatory Logic. Available at: https://crypto.stanford.edu/~blynn/lambda/cl.html.
Lynn, B. (2022) Kiselyov Combinator Translation. Available at: https://crypto.stanford.edu/~blynn/lambda/kiselyov.html.
Mahler, T. (2021) Implementing a Functional Language with Graph Reduction. Available at: https://thma.github.io/posts/2021-12-27-Implementing-a-functional-language-with-Graph-Reduction.html.
Mahler, T. (2021) λ-Calculus, Combinatory Logic and Cartesian Closed Categories. Available at: https://thma.github.io/posts/2021-04-04-Lambda-Calculus-Combinatory-Logic-and-Cartesian-Closed-Categories.html.
Mahler, T. (2022) Evaluating SKI combinators as native Haskell functions. Available at: https://thma.github.io/posts/2022-02-05-Evaluating-SKI-combinators-as-native-Haskell-functions.html.
Naylor, M. (2008) Evaluating Haskell in Haskell, The Monad.Reader. Edited by W. Swierstra, Issue 10.
Seo, K.Y. (2016) Write you an interpreter. Available at: http://kseo.github.io/posts/2016-12-30-write-you-an-interpreter.html.
Turner, D.A. (1979) ‘A new implementation technique for applicative languages’, Software: Practice and Experience, 9(1), pp. 31–49.
Posted in haskell | Tagged , , , , , , | 6 Comments

Competitive Programming in Haskell: two more DP challenges

Continuing the series on dynamic programming, I just have a couple challenge problems for you today. I have indeed solved both of these problems in Haskell, but I don’t yet know how to write elegant solutions! There is a reason that the techniques covered in my previous posts aren’t quite good enough.

Feel free to discuss in the comments! I’m hoping that I can learn some new approaches from some of my readers. I will probably post some hints in the comments towards the right recurrences, so don’t look at the comments if you don’t want any spoilers.

Posted in competitive programming, haskell | Tagged , , , | 7 Comments

Nested folds

I’m finally getting around to reading Algorithm Design with Haskell (hereafter abbreviated as ADH), by Jeremy Gibbons and Richard Bird. I’ve had it for a while, and I have no excuse for waiting this long to read it, but anyway. I’m enjoying it so far, and wanted to share something I (indirectly) learned. I’m sure there are some who already know this, but I didn’t. I’ll share both the fun takeaway and then also the interesting, roundabout path I took to get there.

Composed folds are nested folds

Here’s the punchline:

  • foldl :: (b -> a -> b) -> b -> [a] -> b
  • foldl . foldl :: (b -> a -> b) -> b -> [[a]] -> b
  • foldl . foldl . foldl :: (b -> a -> b) -> b -> [[[a]]] -> b

Actually, it’s a bit more general than this, since foldl works over any Foldable, not just lists: in fact, foldl . foldl . foldl can be used to fold any t1 (t2 (t3 a)) as long as t1, t2, and t3 are all instances of Foldable. For example, here is how we can add up all the integers contained in a Maybe (Tree [Int]):

λ> (foldl . foldl . foldl) (+) 0 (Just (Node [1,2,3] [Node [5,6] [], Node [7] [], Node [9,12] [Node [6] []]]))
51

We can make sense of this if we look at the type of foldl. Its type is

Foldable t => (b -> a -> b) -> b -> t a -> b

and we usually think of it as taking three arguments: a combining function of type b -> a -> b, an initial value of type b, and a structure to fold. But we can also think of it as a one-argument function. That is, it takes a function of type b -> a -> b and transforms it into a function of type b -> t a -> b:

foldl :: Foldable t => (b -> a -> b) -> (b -> t a -> b)

Due to the magic of currying this is equivalent, and the second set of parentheses above is redundant. However, with this change of perspective it is easy to see what’s going on: the result of foldl is a function of type b -> t a -> b, which has the right shape to be the argument of foldl again, but this time with t a in place of a, yielding

(b -> t a -> b) -> (b -> t2 (t a) -> b)

and so on.

What about foldr?

foldr :: Foldable t => (a -> b -> b) -> (b -> t a -> b)

The shapes of the input and output to foldr don’t quite match, but they will if we throw in an extra flip, which switches the arguments of a two-argument function. So we can either do

foldr . flip :: Foldable t => (b -> a -> b) -> (b -> t a -> b)

if we want to match the type of foldl, or

flip . foldr :: Foldable t => (a -> b -> b) -> (t a -> b -> b)

In any case, we can now iterate just as with foldl; for example

foldr . flip . foldr . flip . foldr :: (a -> b -> b) -> (b -> t1 (t2 (t3 a)) -> b)

(with some appropriate Foldable constraints thrown in as well).

My roundabout journey

So how did I come to realize this? Sometimes the journey is more interesting than the destination. The first chapter of ADH talks about the foldr fusion rule, which says that

h . foldr f e == foldr g (h e)

as long as h (f x y) == g x (h y) for all x and y. In other words, if we have a foldr followed by a function h, we can turn this into a single foldr (i.e. “fuse away” the h) as long as we can find an appropriate function g that satisfies the given criterion.

One of the exercises asks us to use foldr fusion to simplify

foldr f e . concat

which performs a fold over a nested list by first flattening the list and then doing a fold. You may wish to go try it yourself before reading on!

The solution

We can compute as follows, where g is some function we will need to define appropriately:

  foldr f e . concat
=                              { definition of concat }
  foldr f e . foldr (++) []
=                              { foldr fusion law, with h = foldr f e }
  foldr g (foldr f e [])
=                              { definition of foldr }
  foldr g e

According to the fusion condition, we need g to satisfy foldr f e (x ++ y) == g x (foldr f e y). This was already given as a previous exercise; I actually solved it by thinking about how foldr works and intuiting the right answer, but we can also calculate it using a second round of foldr fusion:

  g x (foldr f e y)
=
  foldr f e (x ++ y)
=                              { definition of (++) }
  foldr f e (foldr (:) y x)
=                              { foldr fusion }
  foldr f (foldr f e y) x

The last equation follows since foldr f e (a : b) = f a (foldr f e b) by definition of foldr. Hence, using z in place of foldr f e y, we have g x z = foldr f z x = flip (foldr f) x z, and so g = flip (foldr f), and we have

foldr f e . concat = foldr (flip (foldr f)) e

We can simplify even further: if we define nestedFold f e = foldr f e . concat = foldr (flip (foldr f)) e then we can eta-reduce to get nestedFold = foldr . flip . foldr.

When I derived this, my eyes bugged out and I started playing around with it, which is how I ended up figuring out the thing with foldl. Presumably, one could use the similar foldl fusion law to fuse foldl f e . concat and end up deriving the foldl . foldl result; I’ll leave that as an exercise for interested readers.

Posted in haskell | Tagged , , , , | 3 Comments

Dynamic programming in Haskell: automatic memoization

This is part 2 of a promised multi-part series on dynamic programming in Haskell. As a reminder, we’re using Zapis as a sample problem. In this problem, we are given a sequence of opening and closing brackets (parens, square brackets, and curly braces) with question marks, and have to compute the number of different ways in which the question marks could be replaced by brackets to create valid, properly nested bracket sequences.

Last time, we developed some code to efficiently solve this problem using a mutually recursive pair of a function and a lookup table represented by a lazy, immutable array. This solution is pretty good, but it leaves a few things to be desired:

  • It requires defining both a function and a lazy, immutable array, and coming up with names for them.
  • When defining the function, we have to remember to index into the array instead of calling the function recursively, and there is nothing that will warn us if we forget.

An impossible dream

Wouldn’t it be cool if we could just write the recursive function, and then have some generic machinery make it fast for us by automatically generating a memo table?

In other words, we’d like a magic memoization function, with a type something like this:

memo :: (i -> a) -> (i -> a)

Then we could just define our slow, recursive function normally, wave our magic memo wand over it, and get a fast version for free!

This sounds lovely, of course, but there are a few problems:

  • Surely this magic memo function won’t be able to work for any type i. Well, OK, we can add something like an Ix i constraint and/or extra arguments to make sure that values of type i can be used as (or converted to) array indices.

  • How can memo possibly know how big of a table to allocate? One simple way to solve this would be to provide the table size as an extra explicit argument to memo. (In my next post we’ll also explore some clever things we can do when we don’t know in advance how big of a table we will need.)

  • More fundamentally, though, our dream seems impossible: given a function i -> a, the only thing the memo function can do is call it on some input of type i; if the i -> a function is recursive then it will go off and do its recursive thing without ever consulting a memo table, defeating the entire purpose.

… or is it?

For now let’s ignore the fact that our dream seems impossible and think about how we could write memo. The idea is to take the given (i -> a) function and first turn it into a lookup table storing a value of type a for each i; then return a new i -> a function which works by just doing a table lookup.

From my previous post we already have a function to create a table for a given function:

tabulate :: Ix i => (i,i) -> (i -> a) -> Array i a
tabulate rng f = listArray rng (map f $ range rng)

The inverse function, which turns an array back into a function, is just the array indexing operator, with extra parentheses around the i -> a to emphasize the shift in perspective:

(!) :: Ix i => Array i a -> (i -> a)

So we can define memo simply as the composition

memo :: Ix i => (i,i) -> (i -> a) -> (i -> a)
memo rng = (!) . tabulate rng

This is nifty… but as we already saw, it doesn’t help very much… right? For example, let’s define a recursive (slow!) Fibonacci function, and apply memo to it:

{-# LANGUAGE LambdaCase #-}

fib :: Int -> Integer
fib = \case
  0 -> 0
  1 -> 1
  n -> fib (n-1) + fib (n-2)

fib' :: Int -> Integer
fib' = memo (0,1000) fib

As you can see from the following ghci session, calling, say, fib' 35 is still very slow the first time, since it simply calls fib 35 which does its usual exponential recursion. However, if we call fib' 35 a second time, we get the answer instantly:

λ> :set +s
λ> fib' 35
9227465
(4.18 secs, 3,822,432,984 bytes)
λ> fib' 35
9227465
(0.00 secs, 94,104 bytes)

This is better than nothing, but it’s not really the point. We want it to be fast the first time by looking up intermediate results in the memo table. And trying to call fib' on bigger inputs is still going to be completely hopeless.

The punchline

All might seem hopeless at this point, but we actually have everything we need—all we have to do is just stick the call to memo in the definition of fib itself!

fib :: Int -> Integer
fib = memo (0,1000) $ \case
  0 -> 0
  1 -> 1
  n -> fib (n-1) + fib (n-2)

Magically, fib is now fast:

λ> fib 35
9227465
(0.00 secs, 94,096 bytes)
λ> fib 1000
43466557686937456435688527675040625802564660517371780402481729089536555417949051890403879840079255169295922593080322634775209689623239873322471161642996440906533187938298969649928516003704476137795166849228875
(0.01 secs, 807,560 bytes)

This solves all our problems. We only have to write a single definition, which is a directly recursive function, so it’s hard to mess it up. The only thing we have to change is to stick a call to memo (with an appropriate index range) on the front; the whole thing is elegant and short.

How does this even work, though? At first glance, it might seem like it will generate a new table with every recursive call to fib, which would obviously be a disaster. However, that’s not what happens: there is only a single, top-level definition of fib, and it is defined as the function which looks up its input in a certain table. Every time we call fib we are calling that same, unique top-level function which is defined in terms of its (unique, top-level) table. So this ends up being equivalent to our previous solution—there is a mutually recursive pair of a function and a lookup table—but written in a much nicer, more compact way that doesn’t require us to explicitly name the table.

So here’s our final solution for Zapis. As you can see, the extra code we have to write in order to memoize our recurrence boils down to about five lines (two of which are type signatures and could be omitted). This is definitely a technique worth knowing!

{-# LANGUAGE LambdaCase #-}

import Control.Arrow
import Data.Array

main = interact $ lines >>> last >>> solve >>> format

format :: Integer -> String
format = show >>> reverse >>> take 5 >>> reverse

tabulate :: Ix i => (i,i) -> (i -> a) -> Array i a
tabulate rng f = listArray rng (map f $ range rng)

memo :: Ix i => (i,i) -> (i -> a) -> (i -> a)
memo rng = (!) . tabulate rng

solve :: String -> Integer
solve str = c (0,n)
  where
    n = length str
    s = listArray (0,n-1) str

    c :: (Int, Int) -> Integer
    c = memo ((0,0), (n,n)) $ \case
      (i,j)
        | i == j           -> 1
        | even i /= even j -> 0
        | otherwise        -> sum
          [ m (s!i) (s!k) * c (i+1,k) * c (k+1, j)
          | k <- [i+1, i+3 .. j-1]
          ]

m '(' ')'                = 1
m '[' ']'                = 1
m '{' '}'                = 1
m '?' '?'                = 3
m b '?' | b `elem` "([{" = 1
m '?' b | b `elem` ")]}" = 1
m _ _                    = 0
Posted in competitive programming, haskell | Tagged , , , | 4 Comments

Dynamic programming in Haskell: lazy immutable arrays

This is part 1 of a promised multi-part series on dynamic programming in Haskell. As a reminder, we’re using Zapis as a sample problem. In this problem, we are given a sequence of opening and closing brackets (parens, square brackets, and curly braces) with question marks, and have to compute the number of different ways in which the question marks could be replaced by brackets to create valid, properly nested bracket sequences.

Last time, we developed a recurrence for this problem and saw some naive, directly recursive Haskell code for computing it. Although this naive version is technically correct, it is much too slow, so our goal is to implement it more efficiently.

Mutable arrays?

Someone coming from an imperative background might immediately reach for some kind of mutable array, e.g. STUArray. Every time we call the function, we check whether the corresponding array index has already been filled in. If so, we simply return the stored value; if not, we compute the value recursively, and then fill in the array before returning it.

This would work, but there is a better way!

Immutable arrays

While mutable arrays occasionally have their place, we can surprisingly often get away with immutable arrays, where we completely define the array up front and then only use it for fast lookups afterwards.

  • If the type of the array elements is suitable, and we can initialize the array elements all at once from a list using some kind of formula, map, scan, etc., we should use UArray since it is much faster than Array.
  • However, UArray is strict in the elements, and the elements must be of a type that can be stored unboxed. If we need a more complex element type, or we need to compute the array recursively (where some elements depend on other elements), we can use Array.

What about the vector library, you ask? Well, it’s a very nice library, and quite fast, but unfortunately it is not available on many judging platforms, so I tend to stick to array to be safe. However, if you’re doing something like Advent of Code or Project Euler where you get to run the code on your own machine, then you should definitely reach for vector.

Lazy, recursive, immutable arrays

In my previous post on topsort we already saw the basic idea: since Arrays are lazy in their elements, we can define them recursively; the Haskell runtime then takes care of computing the elements in a suitable order. Previously, we saw this applied to automatically compute a topological sort, but more generally, we can use it to fill out a table of values for any recurrence.

So, as a first attempt, let’s just replace our recursive c function from last time with an array. I’ll only show the solve function for now; the rest of the code remains the same. (Spoiler alert: this solution works, but it’s ugly. We’ll develop much better solutions later.)

solve :: String -> Integer
solve str = c!(0,n)
  where
    n = length str
    s = listArray (0,n-1) str

    c :: Array (Int, Int) Integer
    c = array ((0,0),(n,n)) $
      [ ((i,i), 1) | i <- [0..n] ]
      ++
      [ ((i,j),0) | i <- [0..n], j <- [0..n], even i /= even j ]
      ++
      [ ((i,j),v)
      | i <- [0..n], j <- [0..n], i /= j, even i == even j
      , let v = sum [ m (s!i) (s!k) * c!(i+1,k) * c!(k+1,j) | k <- [i+1, i+3 .. j-1]]
      ]

We use the array function to create an array, which takes first a pair of indices specifying the index range, and then a list of (index, value) pairs. (The listArray function can also be particularly useful, when we have a list of values which are already in index order, as in the definition of s.)

This solution is accepted, and it’s quite fast (0.04s for me). However, it’s really ugly, and although it’s conceptually close to our directly recursive function from before, the code is almost unrecognizably different. It’s ugly that we have to repeat conditions like i /= j and even i == even j, and binders like i <- [0..n]; the multiple list comprehensions and nested pairs like ((i,j),v) are kind of ugly, and the fact that this is implementing a recurrence is completely obscured.

However, I included this solution as a first step because for a long time, after I learned about using lazy immutable arrays to implement dynamic programming in Haskell, this was the kind of solution I wrote! Indeed, if you just think about the idea of creating a recursively defined array, this might be the kind of thing you come up with: we define an array c using the array function, then we have to list all its elements, and we get to refer to c along the way.

Mutual recursion to the rescue

Most of the ugliness comes from losing sight of the fact that there is a function mapping indices to values: we simply listed out all the function’s input/output pairs without getting to use any of Haskell’s very nice facilities for defining functions! So we can clean up the code considerably if we make a mutually recursive pair of an array and a function: the array values are defined using the function, and the function definition can look up values in the array.

solve :: String -> Integer
solve str = cA!(0,n)
  where
    n = length str
    s = listArray (0,n-1) str

    cA :: Array (Int, Int) Integer
    cA = array ((0,0),(n,n)) $
      [ ((i,j), c (i,j)) | i <- [0 .. n], j <- [0 .. n] ]

    c :: (Int, Int) -> Integer
    c (i,j)
      | i == j           = 1
      | even i /= even j = 0
      | otherwise        = sum
        [ m (s!i) (s!k) * cA ! (i+1,k) * cA ! (k+1,j)
        | k <- [i+1, i+3 .. j-1]
        ]

Much better! The c function looks much the same as our naive version from before, with the one difference that instead of calling itself recursively, it looks up values in the array cA. The array, in turn, is simply defined as a lookup table for the outputs of the function.

Generalized tabulation

One nice trick we can use to simplify the code a bit more is to use the range function to generate the list of all valid array indices, and then just map the c function over this. This also allows us to use the listArray function, since we know that the range will generate the indices in the right order.

cA :: Array (Int, Int) Integer
cA = listArray rng $ map c (range rng)
  where
    rng = ((0,0), (n,n))

In fact, we can abstract this into a useful little function to create a lookup table for a function:

tabulate :: Ix i => (i,i) -> (i -> a) -> Array i a
tabulate rng f = listArray rng (map f $ range rng)

(We can generalize this even more to make it work for UArray as well as Array, but I’ll stop here for now. And yes, I intentionally named this to echo the tabulate function from the adjunctions package; Array i is indeed a representable functor, though it’s not really possible to express without dependent types.)

The solution so far

Putting it all together, here’s our complete solution so far. It’s pretty good, and in fact it’s organized in a very similar way to Soumik Sarkar’s dynamic programming solution to Chemist’s Vows. (However, there’s an even better solution coming in my next post!)

import Control.Arrow
import Data.Array

main = interact $ lines >>> last >>> solve >>> format

format :: Integer -> String
format = show >>> reverse >>> take 5 >>> reverse

tabulate :: Ix i => (i,i) -> (i -> a) -> Array i a
tabulate rng f = listArray rng (map f $ range rng)

solve :: String -> Integer
solve str = cA!(0,n)
  where
    n = length str
    s = listArray (0,n-1) str

    cA :: Array (Int, Int) Integer
    cA = tabulate ((0,0),(n,n)) c

    c :: (Int, Int) -> Integer
    c (i,j)
      | i == j           = 1
      | even i /= even j = 0
      | otherwise        = sum
        [ m (s!i) (s!k) * cA ! (i+1,k) * cA ! (k+1,j)
        | k <- [i+1, i+3 .. j-1]
        ]

m '(' ')'                = 1
m '[' ']'                = 1
m '{' '}'                = 1
m '?' '?'                = 3
m b '?' | b `elem` "([{" = 1
m '?' b | b `elem` ")]}" = 1
m _ _                    = 0

Coming up next: automatic memoization!

So what’s not to like about this solution? Well, I still don’t like the fact that we have to define a mutually recursive array and function. Conceptually, I want to name them both c (or whatever) since they are really isomorphic representations of the exact same mathematical function. It’s annoying that I have to make up a name like cA or c' or whatever for one of them. I also don’t like that we have to remember to do array lookups instead of recursive calls in the function—and if we forget, Haskell will not complain! It will just be really slow.

Next time, we’ll see how to use some clever ideas from Conal Elliot’s MemoTrie package (which themselves ultimately came from a paper by Ralf Hinze) to solve these remaining issues and end up with some really beautiful code!

Posted in competitive programming, haskell | Tagged , , , | 2 Comments

Competitive programming in Haskell: introduction to dynamic programming

In my previous post, I challenged you to solve Zapis. In this problem, we are given a sequence of opening and closing brackets (parens, square brackets, and curly braces) with question marks, and have to compute the number of different ways in which the question marks could be replaced by brackets to create valid, properly nested bracket sequences.

For example, given (??), the answer is 4: we could replace the question marks with any matched pair (either (), [], or {}), or we could replace them with )(, resulting in ()().

An annoying aside

One very annoying thing to mention about this problem is that it requires us to output the last 5 digits of the answer. At first, I interpreted that to mean “output the answer modulo 10^5”, which would be a standard sort of condition for a combinatorics problem, but that’s not quite the same thing, in a very annoying way: for example, if the answer is 2, we are supposed to output 2; but if the answer is 1000000002, we are supposed to output 00002, not 2! So simply computing the answer modulo 10^5 is not good enough; if we get a final answer of 2, we don’t know whether we are supposed to pad it with zeros. I could imagine keeping track of both the result modulo 10^5 along with a Boolean flag telling us whether the number has ever overflowed; we have to pad with zeros iff the flag is set at the end. I’m pretty sure this would work. But for this problem, it turns out that the final answer is at most “only” about 100 digits, so we can just compute the answer exactly as an Integer and then literally show the last 5 digits.

A recurrence

Now, how to compute the answer? For this kind of problem the first step is to come up with a recurrence. Let s[0 \dots n-1] be the given string, and let c(i,j) be the number of ways to turn the substring s[i \dots j-1] into a properly nested sequence of brackets, so ultimately we want to compute the value of c(0,n). (Note we make c(i,j) correspond to the substring which includes i but excludes j, which means, for example, that the length of the substring is j-i.) First, some base cases:

  • c(i,i) = 1 since the empty string always counts as properly nested.
  • c(i,j) = 0 if i and j have different parity, since any properly nested string must have even length.

Otherwise, s[i] had better be an opening bracket of some kind, and we can try matching it with each of s[i+1], s[i+3], s[i+5], …, s[j-1]. In general, matching s[i] with s[k] can be done in either 0, 1, or 3 ways depending on whether they are proper opening and closing brackets and whether any question marks are involved; then we have c(i+1,k) ways to make the substring between s[i] and s[k] properly nested, and c(k+1,j) ways for the rest of the string following s[k]. These are all independent, so we multiply them. Overall, we get this:

c(i,j) = \begin{cases} 1 & i = j \\ 0 & i \not \equiv j \pmod 2 \\ \displaystyle \sum_{k \in [i+1, i+3, \dots, j-1]} m(s[i], s[k]) \cdot c(i+1,k) \cdot c(k+1,j) & \text{otherwise} \end{cases}

where m(x,y) counts the number of ways to make x and y into a matching pair of brackets: it returns 0 if the two characters cannot possibly be a matching open-close pair (either because they do not match or because one of them is the wrong way around); 1 if they match, and at most one of them is a question mark; and 3 if both are question marks.

How do we come up with such recurrences in the first place? Unfortunately, Haskell doesn’t really make this any easier—it requires some experience and insight. However, what we can say is that Haskell makes it very easy to directly code a recurrence as a recursive function, to play with it and ensure that it gives correct results for small input values.

A naive solution

To that end, if we directly code up our recurrence in Haskell, we get the following naive solution:

import Control.Arrow
import Data.Array

main = interact $ lines >>> last >>> solve >>> format

format :: Integer -> String
format = show >>> reverse >>> take 5 >>> reverse

solve :: String -> Integer
solve str = c (0,n)
  where
    n = length str
    s = listArray (0,n-1) str

    c :: (Int, Int) -> Integer
    c (i,j)
      | i == j           = 1
      | even i /= even j = 0
      | otherwise        = sum
        [ m (s!i) (s!k) * c (i+1,k) * c (k+1,j)
        | k <- [i+1, i+3 .. j-1]
        ]

m '(' ')'                = 1
m '[' ']'                = 1
m '{' '}'                = 1
m '?' '?'                = 3
m b '?' | b `elem` "([{" = 1
m '?' b | b `elem` ")]}" = 1
m _ _                    = 0

This solution is correct, but much too slow—it passes the first four test cases but then fails with a Time Limit Exceeded error. In fact, it takes exponential time in the length of the input string, because it has a classic case of overlapping subproblems. Our goal is to compute the same function, but in a way that is actually efficient.

Dynamic programming, aka memoizing recurrences

I hate the name “dynamic programming”—it conveys zero information about the thing that it names, and was essentially invented as a marketing gimmick. Dynamic programming is really just memoizing recurrences in order to compute them more efficiently. By memoizing we mean caching some kind of mapping from input to output values, so that we only have to compute a function once for each given input value; on subsequent calls with a repeated input we can just look up the corresponding output. There are many, many variations on the theme, but memoizing recurrences is really the heart of it.

In imperative languages, dynamic programming is often carried out by filling in tables via nested loops—the fact that there is a recurrence involved is obscured by the implementation. However, in Haskell, our goal will be to write code that is as close as possible to the above naive recursive version, but still actually efficient. Over the next few posts we will discuss several techniques for doing just that.

  • In part 1, we will explore the basic idea of using lazy, recursive, immutable arrays (which we have already seen in a previous post).
  • In part 2, we will use ideas from Conal Elliot’s MemoTrie package (and ultimately from a paper by Ralf Hinze) to clean up the code and make it a lot closer to the naive version.
  • This post contains a couple challenge problems that can’t quite be solved using the techniques in the previous posts.
  • At some point perhaps we’ll discuss how to memoize functions with infinite (or just very large) domains.
  • There may very well end up being more parts… we’ll see where it ends up!

Along the way I’ll also drop more links to relevant background. This will ultimately end up as a chapter in the book I’m slowly writing, and I’d like to make it into the definitive reference on dynamic programming in Haskell—so any thoughts, comments, links, etc. are most welcome!

Posted in competitive programming, haskell | Tagged , , , | 8 Comments

Competitive programming in Haskell: parsing with an NFA

In my previous post, I challenged you to solve Chemist’s Vows. In this problem, we have to decide which words can be made by concatenating atomic element symbols. So this is another parsing problem; but unlike the previous problem, element symbols are not prefix-free. For example, B and Be are both element symbols. So, if we see BE..., we don’t immediately know whether we should parse it as Be, or as B followed by an element that starts with E (such as Er).

A first try

A parsing problem, eh? Haskell actually shines in this area because of its nice parser combinator libraries. The Kattis environment does in fact have the parsec package available; and even on platforms that don’t have parsec, we can always use the Text.ParserCombinators.ReadP module that comes in base. So let’s try throwing one of those packages at the problem and see what happens!

If we try using parsec, we immediately run into problems; honestly, I don’t even know how to solve the problem using parsec. The problem is that <|> represents left-biased choice. If we parse p1 <|> p2 and parser p1 succeeds, then we will never consider p2. But for this parsing problem, because the symbols are not prefix-free, sometimes we can’t know which of two options we should have picked until later.

ReadP, on the other hand, explicitly has both biased and unbiased choice operators, and can return a list of possible parses instead of just a single parse. That sounds promising! Here’s a simple attempt using ReadP: to parse a single element, we use an unbiased choice over all the element names; then we use many parseElement <* eof to parse each word, and check whether there are any successful parses at all.

{-# LANGUAGE OverloadedStrings #-}

import           Control.Arrow
import           Data.Bool
import qualified Data.ByteString.Lazy.Char8   as C
import           Text.ParserCombinators.ReadP (ReadP, choice, eof, many,
                                               readP_to_S, string)

main = C.interact $
  C.lines >>> drop 1 >>> map (solve >>> bool "NO" "YES") >>> C.unlines

solve :: C.ByteString -> Bool
solve s = case readP_to_S (many parseElement <* eof) (C.unpack s) of
  [] -> False
  _  -> True

elements :: [String]
elements = words $
  "h he li be b c n o f ne na mg al si p s cl ar k ca sc ti v cr mn fe co ni cu zn ga ge as se br kr rb sr y zr nb mo tc ru rh pd ag cd in sn sb te i xe cs ba hf ta w re os ir pt au hg tl pb bi po at rn fr ra rf db sg bh hs mt ds rg cn fl lv la ce pr nd pm sm eu gd tb dy ho er tm yb lu ac th pa u np pu am cm bk cf es fm md no lr"

parseElement :: ReadP String
parseElement = choice (map string elements)

Unfortunately, this fails with a Time Limit Exceeded error (it takes longer than the allotted 5 seconds). The problem is that backtracking and trying every possible parse like this is super inefficient. One of the secret test inputs is almost cerainly constructed so that there are an exponential number of ways to parse some prefix of the input, but no way to parse the entire thing. As a simple example, the string crf can be parsed as either c rf (carbon + rutherfordium) or cr f (chromium + fluorine), so by repeating crf n times we can make a string of length 3n which has 2^n different parses. If we fed this string to the ReadP solution above, it would quickly succeed with more or less the first thing that it tried. However, if we stick a letter on the end that does not occur in any element symbol (such as q), the result will be an unparseable string, and the ReadP solution will spend a very long time backtracking through exponentially many parses that all ultimately fail.

Solution

The key insight is that we don’t really care about all the different possible parses; we only care whether the given string is parseable at all. At any given point in the string, there are only two possible states we could be in: we could be finished reading one element symbol and about to start reading the next one, or we could be in the middle of reading a two-letter element symbol. We can just scan through the string and keep track of the set of (at most two) possible states; in other words, we will simulate an NFA which accepts the language of strings composed of element symbols.

First, some setup as before.

{-# LANGUAGE OverloadedStrings #-}

import           Control.Arrow              ((>>>))
import           Data.Array                 (Array, accumArray, (!))
import           Data.Bool                  (bool)
import qualified Data.ByteString.Lazy.Char8 as C
import           Data.List                  (partition, nub)
import           Data.Set                   (Set)
import qualified Data.Set                   as S

main = C.interact $
  C.lines >>> drop 1 >>> map (solve >>> bool "NO" "YES") >>> C.unlines

elements :: [String]
elements = words $
  "h he li be b c n o f ne na mg al si p s cl ar k ca sc ti v cr mn
fe co ni cu zn ga ge as se br kr rb sr y zr nb mo tc ru rh pd ag cd
in sn sb te i xe cs ba hf ta w re os ir pt au hg tl pb bi po at rn
fr ra rf db sg bh hs mt ds rg cn fl lv la ce pr nd pm sm eu gd tb dy
ho er tm yb lu ac th pa u np pu am cm bk cf es fm md no lr"

Now, let’s split the element symbols into one-letter and two-letter symbols:

singles, doubles :: [String]
(singles, doubles) = partition ((==1).length) elements

We can now make boolean lookup arrays that tell us whether a given letter occurs as a single-letter element symbol (single) and whether a given letter occurs as the first letter of a two-letter symbol (lead). We also make a Set of all two-letter element symbols, for fast lookup.

mkAlphaArray :: [Char] -> Array Char Bool
mkAlphaArray cs = accumArray (||) False ('a', 'z') (zip cs (repeat True))

single, lead :: Array Char Bool
[single, lead] = map (mkAlphaArray . map head) [singles, doubles]

doubleSet :: Set String
doubleSet = S.fromList doubles

Now for simulating the NFA itself. There are two states we can be in: START means we are about to start and/or have just finished reading an element symbol; SEEN c means we have seen the first character of some element (c) and are waiting to see another.

data State = START | SEEN Char
  deriving (Eq, Ord, Show)

Our transition function takes a character c and a state and returns a set of all possible next states (we just use a list since these sets will be very small). If we are in the START state, we could end up in the START state again if c is a single-letter element symbol; we could also end up in the SEEN c state if c is the first letter of any two-letter element symbol. On the other hand, if we are in the SEEN x state, then we have to check whether xc is a valid element symbol; if so, we return to START.

delta :: Char -> State -> [State]
delta c START    = [START | single!c] ++ [SEEN c | lead!c]
delta c (SEEN x) = [START | [x,c] `S.member` doubleSet]

We can now extend delta to act on a set of states, giving us the set of all possible resulting states; the drive function then iterates this one-letter transition over an entire input string. Finally, to solve the problem, we start with the singleton set [START], call drive using the input string, and check whether START (which is also the only accepting state) is an element of the resulting set of states.

trans :: Char -> [State] -> [State]
trans c sts = nub (sts >>= delta c)

drive :: C.ByteString -> ([State] -> [State])
drive = C.foldr (\c -> (trans c >>>)) id

solve :: C.ByteString -> Bool
solve s = START `elem` drive s [START]

And that’s it! This solution is accepted in 0.27 seconds (out of a maximum allowed 5 seconds).

For next time

  • If you want to practice the concepts from my past couple posts, give Haiku a try.
  • For my next post, I challenge you to solve Zapis!
Posted in competitive programming, haskell | Tagged , , , , | 12 Comments

New ko-fi page: help me attend ICFP!

tl;dr: if you appreciate my past or ongoing contributions to the Haskell community, please consider helping me get to ICFP by donating via my new ko-fi page!

ko-fi

Working at a small liberal arts institution has some tremendous benefits (close interaction with motivated students, freedom to pursue the projects I want rather than jump through a bunch of hoops to get tenure, fantastic colleagues), and I love my job. But there are also downsides; the biggest ones for me are the difficulty of securing enough travel funding, and, relatedly, the difficulty of cultivating and maintaining collaborations.

I would really like to be able to attend ICFP in Seattle this September; the last time I was able to attend ICFP in person was 2019 in Berlin. With transportation, lodging, food, and registration fees, it will probably come to about $3000. I can get a grant from my instutition to pay for up to $1200, but that still leaves a big gap.

As I was brainstorming other sources of funding, it dawned on me that there are probably many people who have been positively impacted by my contributions to the Haskell community (e.g. CIS 194, the Typeclassopedia, diagrams, split, MonadRandom, burrito metaphors…) and/or would like to support my ongoing work (competitive programming in Haskell, swarm, disco, ongoing package maintenance…) and would be happy to donate a bit.

So, to that end, I have set up a ko-fi page.

  • If you have been positively impacted by my contributions and would like to help me get to ICFP this fall, one-time donations — even very small amounts — are greatly appreciated! I’m not going to promise any particular rewards, but if you’re at ICFP I will definitely find you to say thanks!

  • Thinking beyond this fall, ideally this could even become a reliable source of funding to help me travel to ICFP or other collaboration opportunities every year. To that end, if you’re willing to sign up for a recurring monthly donation, that would be amazing — think of it as supporting my ongoing work: blog posts (and book) on competitive programming in Haskell, Swarm and Disco development, and ongoing package maintenance. I will post updates on ko-fi with things I’m thinking about and working on; I am also going to try to publish more frequent blog posts, at least in the near term.

Thank you, friends — I hope to see many people in Seattle! Next up: back to your regularly scheduled competitive programming!

Posted in haskell, meta | Tagged , , , , , | Leave a comment

Competitive programming in Haskell: tries

In my previous post, I challenged you to solve Alien Math, which is about reading numbers in some base B, but with a twist. We are given a list of B strings representing the names of the digits 0 through B-1, and a single string describing a number, consisting of concatenated digit names. For example, if B = 3 and the names of the digits are zero, one, two, then we might be given a string like twotwozerotwoone, which we should interpret as 22021_3 = 223_{10}. Crucially, we are also told that the digit names are prefix-free, that is, no digit name is a prefix of any other. But other than that, the digit names could be really weird: they could be very different lengths, some digit names could occur as substrings (just not prefixes) of others, digit names could share common prefixes, and so on. So this is really more of a parsing problem than a math problem; once we have parsed the string as a list of digits, converting from base B is the easy part.

One simple way we can do this is to define a map from digit names to digits, and simply look up each prefix of the given string until we find a hit, then chop off that prefix and start looking at successive prefixes of the remainder. This takes something like O(n^2 \lg n) time in the worst case (I think)—but this is actually fine since n is at most 300. This solution is accepted and runs in 0.00 seconds for me.

Tries

However, I want to talk about a more sophisticated solution that has better asymptotic time complexity and generalizes nicely to other problems. Reading a sequence of strings from a prefix-free set should make you think of Huffman coding, if you’ve ever seen that before. In general, the idea is to define a trie containing all the digit names, with each leaf storing the corresponding digit. We can then scan through the input one character at a time, keeping track of our current position in trie, and emit a digit (and restart at the root) every time we reach a leaf. This should run in O(n) time.

Let’s see some generic Haskell code for tries (this code can also be found at byorgey/comprog-hs/Trie.hs on GitHub). First, some imports, a data type definition, and emptyTrie and foldTrie for convenience:

module Trie where

import           Control.Monad              ((>=>))
import qualified Data.ByteString.Lazy.Char8 as C
import           Data.List                  (foldl')
import           Data.Map                   (Map, (!))
import qualified Data.Map                   as M
import           Data.Maybe                 (fromMaybe)

data Trie a = Trie
  { trieSize :: !Int
  , value    :: !(Maybe a)
  , children :: !(Map Char (Trie a))
  }
  deriving Show

emptyTrie :: Trie a
emptyTrie = Trie 0 Nothing M.empty

-- | Fold a trie into a summary value.
foldTrie :: (Int -> Maybe a -> Map Char r -> r) -> Trie a -> r
foldTrie f (Trie n b m) = f n b (M.map (foldTrie f) m)

A trie has a cached size (we could easily generalize this to store any sort of monoidal annotation), a possible value (i.e. the value associated with the empty string key, if any), and a map from characters to child tries. The cached size is not needed for this problem, but is included since I needed it for some other problems.

Now for inserting a key/value pair into a Trie. This code honestly took me a while to get right! We fold over the given string key, producing for each key suffix a function which will insert that key suffix into a trie. We have to be careful to correctly update the size, which depends on whether the key being inserted already exists—so the recursive go function actually returns a pair of a new Trie and an Int representing the change in size.

-- | Insert a new key/value pair into a trie, updating the size
--   appropriately.
insert :: C.ByteString -> a -> Trie a -> Trie a
insert w a t = fst (go w t)
  where
    go = C.foldr
      (\c insSuffix (Trie n v m) ->
         let (t', ds) = insSuffix (fromMaybe emptyTrie (M.lookup c m))
         in  (Trie (n+ds) v (M.insert c t' m), ds)
      )
      (\(Trie n v m) ->
         let ds = if isJust v then 0 else 1
         in  (Trie (n+ds) (Just a) m, ds)
      )

Now we can create an entire Trie in one go by folding over a list of key/value pairs with insert:

-- | Create an initial trie from a list of key/value pairs.  If there
--   are multiple pairs with the same key, later pairs override
--   earlier ones.
mkTrie :: [(C.ByteString, a)] -> Trie a
mkTrie = foldl' (flip (uncurry insert)) emptyTrie

A few lookup functions: one to look up a single character and return the corresponding child trie, and then on top of that we can build one to look up the value associated to an entire string key.

-- | Look up a single character in a trie, returning the corresponding
--   child trie (if any).
lookup1 :: Char -> Trie a -> Maybe (Trie a)
lookup1 c = M.lookup c . children

-- | Look up a string key in a trie, returning the corresponding value
--   (if any).
lookup :: C.ByteString -> Trie a -> Maybe a
lookup = C.foldr ((>=>) . lookup1) value

Finally, a function that often comes in handy for using a trie to decode a prefix-free code. It takes an input string and looks it up character by character; every time it encounters a key which exists in the trie, it emits the corresponding value and then starts over at the root of the trie.

decode :: Trie a -> C.ByteString -> [a]
decode t = reverse . snd . C.foldl' step (t, [])
  where
    step (s, as) c =
      let Just s' = lookup1 c s
      in  maybe (s', as) (\a -> (t, a:as)) (value s')

These tries are limited to string keys, since that is most useful in a competitive programming context, but it is of course possible to make much more general sorts of tries — see Hinze, Generalizing Generalized Tries.

Solution

Finally, we can use our generic tries to solve the problem: read the input, build a trie mapping digit names to values, use the decode function to read the given number, and finally interpret the resulting list of digits in the given base.

import Control.Arrow ((>>>))
import ScannerBS
import Trie

main = C.interact $ runScanner tc >>> solve >>> showB

data TC = TC { base :: Integer, digits :: [C.ByteString], number :: C.ByteString }
  deriving (Eq, Show)

tc :: Scanner TC
tc = do
  base <- integer
  TC base <$> (fromIntegral base >< str) <*> str

solve :: TC -> Integer
solve TC{..} = foldl' (\n d -> n*base + d) 0 (decode t number)
  where
    t = mkTrie (zip digits [0 :: Integer ..])

Practice problems

Here are a few other problems where you can profitably make use of tries. Some of these can be solved directly using the Trie code given above; others may require some modifications or enhancements to the basic concept.

For next time

For next time, I challenge you to solve Chemist’s vows!

Posted in competitive programming, haskell | Tagged , , , | 2 Comments