On Linear Scopes (improving their ergonomics without requiring linear constraints)

[Qqwy]

I've been thinking about https://www.tweag.io/blog/2023-03-23-linear-constraints-linearly/ and the related #130 these last few days.

I think that having a Linearly typeclass is a very cool idea. However, for obvious reasons it will probably take quite a while before it might be introduced into GHC and can be used in practical code. So I was thinking about alternative approaches which might be 'good enough' today.

---

1. Removing Ur from alloc

Today in linear-base, the following is the main API entrypoint to work with mutable arrays:

-- In the article, alloc is known as `newMArray`
-- For some other collections, a very similar function `fromList` is used
alloc :: Int -> a -> (MArray a %1 -> Ur b) %1 -> Ur b

(And similar functions exist to work with mutable vectors, mutable sets, mutable hashmaps, etc.)

It is mentioned in the article that the Ur is absolutely necessary here.
But I'm not so sure: The main point as far as I can see is to make sure that the MArray is not allowed to remain after the callback returns. Another way would be:

alloc' :: Immutable b => Int -> a -> (Marray a %1 -> b) %1 -> b

Where Immutable is an (empty) typeclass which is implemented for most types but explicitly not for 'in-place mutation' types.

Indeed, a blanket instance Movable a => Immutable a could be introduced which essentially recovers the original API but without having to wrap/unwrap Ur.
And there might be other types (ex: resource handles which we use linearly to make sure they are cleaned up early and exactly once) that might implement Immutable but are intentionally not Movable.

2. Resolving the 'nˆ2 beside implementations' problem

In the article it is mentioned that because scopes introduced with alloc are 'sticky'. As solution, the following function is shown:

-- In the article, allocBeside is known as `newMArrayBeside`
allocBeside :: Int -> a -> Array b %1 -> (Array a, Array b)

But the problem which is immediately raised, is that if you want to mix and match different kinds of mutable collections, you'd need n^2 functions of this shape (one for each combination).

But what about introducing a typeclass?

-- Side note: This example uses a multi param typeclass with functional constraints
-- but maybe an associated type- or data family might be even more ergonomic to use?
class Allocable a opts | opts -> a where
    -- | Bulds an `a` from the given options.
    --
    -- Safety: This function allows you to obtain an `a` outside of a linear context.
    -- The caller needs to make sure that the resulting `a` is used exactly once!
    unsafeAlloc :: opts %1 -> a

alloc :: (Immutable b, Allocable a opts) => opts %1 -> (a %1 -> b) %1 -> b
alloc opts f = f (unsafeAlloc opts)

allocBeside :: (Allocable a aOpts, Allocable b bOpts) => bOpts -> a %1 -> (a, b)
allocBeside opts a = (a, unsafeAlloc opts)

The implementation for arrays could look like:

data ArrayOptions a  where
  FromElement :: Int -> a -> ArrayOptions a
  FromList :: [a] -> ArrayOptions a -- <- Maybe we can construct this linearly by the way?

instance Allocable (Array a) (ArrayOpts a) where
    unsafeAlloc (FromElement size elem) = Array (unsafeArrAlloc size elem)
    unsafeAlloc (FromList list) = 
        unsafeAlloc (FromElement (Prelude.length list) (error "invariant violation: unintialized array position" :: a))
        & doWrites (Prelude.zip list [0..])
            where
                doWrites :: [(a, Int)] -> Array a %1 -> Array a
                doWrites [] arr = arr
                doWrites ((a, ix) : xs) arr = doWrites xs (unsafeSet ix a arr)

This way the n^2 problem is averted.

---

I am still quite new to Linear Haskell so there's a reasonable chance there's a fundamental issue with this approach which I missed, but at least at first glance it seems promising 😅

What do you think?

[aspiwack]

Hi, just a word to say that I'm traveling for a little while still, and I don't have time to give you're posts the time they deserve. I very much appreciate all the ideas that you're coming up with. I'll get back to you in a week or two.

[Qqwy]

Have a great trip! 😊

[Qqwy]

I realized that we can go a step further in (2 'resolving the n^2 beside implementations problem'). Rather than using a complicated typeclass, we can use a 'marker' class and a few combinator functions:

{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE RankNTypes #-}
module Unique (
    -- | = Class and combinator functions to work with unique values
    Unique,
    scoped,
    consuming,
    -- | = UniqueUnit: the zero-info unique witness type
    UniqueUnit, 
    unit,
) where
import Prelude.Linear

-- | Implemented by types which can used only inside a scoped unique context.
--
-- Instances of this class can be used as 'uniqueness witnesses' to create more unique values.
--
-- = Invariants:
--
-- - A type can either implement @Unique@ or @Movable@ but never both.
-- - A `Unique` has no (safe) normal public constructor. Instead, they are created by type-specific 'unique constructor' functions.
--
class Unique u where

-- | Combinator which allows you to use a unique constructor function to kickstart a unique scope.
--
scoped :: (Unique u, Movable r) => 
    (forall witness. Unique witness => witness %1 -> (u, witness)) -- ^ constructor function which when given a uniqueness witness creates a new unique value
    %1 -> (u %1 -> r)  -- ^ Scoped callback function
    %1 -> r -- ^ Final, non-unique result.
scoped constructor callback = 
    UniqueUnit
    & consuming constructor
    & callback


-- | Combinator which turns a normal unique constructor function into a consuming constructor function: 
-- Given a `b` that you no longer need, consume it after using it as a uniqueness witness to create `a`.
consuming :: (Unique u, Unique witness, Consumable witness)  => (witness %1 -> (u, witness)) %1 -> witness %1 -> u
consuming f = (\(a, b) -> b `lseq` a) . f


-- | The zero-info unique unit type
-- This type is similar to @()@ except that the only way to obtain one is inside a scoped unique context.
--
-- (Similar to the `Source` token in https://github.com/tweag/linear-base/issues/130 but more ergonomic to use)
data UniqueUnit = UniqueUnit

instance Unique UniqueUnit where

instance Consumable UniqueUnit where
    consume UniqueUnit = ()

instance Dupable UniqueUnit where
    dup2 UniqueUnit = (UniqueUnit, UniqueUnit)

-- | Unique constructor for `UniqueUnit`.
unit :: Unique witness => witness %1 -> (UniqueUnit, witness)
unit x = (UniqueUnit, x)

The core idea behind this concept is the structure of 'unique constructor' functions.

A normal constructors has the type ... -> a.
A 'unique constuctor' on the other hand has the type Unique witness => ... -> witness %1 -> (a, witness).
In other words: Only if you already have a Unique value witness, you can use it as to create an a.
Because linearity is 'infectious', the uniqueness of the witness is 'inherited' by the constructed value.
This is how we can ensure that the constructed value can never be accidentally aliased.

As the observant reader might already have realized: Array.allocBeside and the other *Beside functions currently existing in linear-base are essentially such 'unique constructors'. So let's generalize them to work with any Unique:

module Array (...) where

fromList :: (Unique witness) => [a] %1 -> witness %1 -> (Array a, witness)

alloc :: (Unique witness) => Int %1 -> a -> witness %1 -> (Array a, witness)

The neat thing now is that we can recover the other variants of the functions:

fromListScoped :: Movable r => [a] -> (Array a %1 -> r) %1 -> r
fromListScoped list = Unique.scoped (Array.fromList list)

fromListConsuming :: (Unique witness, Consumable witness) => [a] -> witness %1 -> Array a
fromListConsuming list = Unique.consuming (Array.fromList list)

allocScoped :: Movable r => Int -> a -> (Array a %1 -> r) %1 -> r
allocScoped size elem = Unique.scoped (Array.alloc size elem)

allocConsuming :: (Unique witness, Consumable witness) => Int -> a -> witness %1 -> Array a
allocConsuming = Unique.consuming (Array.alloc size elem) 

-- essentially 'withSource' from https://github.com/tweag/linear-base/issues/130
unitScoped :: Movable r => (UniqueUnit %1 -> r) %1 -> r
unitScoped = Unique.scoped Unique.unit

unitConsuming :: (Unique witness, Consumable witness) => witness %1 -> UniqueUnit
unitConsuming = Unique.consuming Unique.unit

And thus for any new mutable collection we only need to provide:

• an instance of the marker typeclass Unique.
• one or more 'unique constructor' functions.

and you get all the other variants and combinations for free!

Here is an example of what you can do with these building blocks today:
_{(This would be even more ergonomic once let bindings can work linearly)}

-- based on the example motivating newMArrayBeside from https://www.tweag.io/blog/2023-03-23-linear-constraints-linearly/
blogExample :: [Int] -> (Int, [Int])
blogExample list =
  Unique.scoped (Array.fromList list) $ \arr ->
    Array.alloc 2 0 arr & \(prefix, arr) ->
    Array.get 0 arr & \(a0, arr) ->
    Array.get 1 arr & \(a1, arr) ->
    ( prefix & Array.set 0 a0 & Array.set 1 a1 & Array.consumingSum, Array.toList arr)

I'll be working a little bit on a small proof-of-concept to further flesh out this idea.

[aspiwack]

Interesting. But something is still missing here. There doesn't seem to be a way to prove that a type is Unique.

Unless the intention is that I can prove Unique at every type, but then this is really unsafe. I think Unique should require some method. Any thought?

[Qqwy]

Since Unique is intended to be a witness to the absence of a Movable instance, I cannot think of any method that would make sense, though maybe there is?

At first I also thought that it would require perfect compliance by the user to not mess up (e.g. add a 'law' and hope that the user would not by accident implement both Movable and Unique for the same type). But maybe it is possible to enforce type safety by adding an erroring instance:

import Data.Kind (Constraint)
import GHC.TypeLits (ErrorMessage (Text), TypeError)

instance (Movable m, MovableTypesAreNeverUnique) => Unique m

type family MovableTypesAreNeverUnique :: Constraint where
  MovableTypesAreNeverUnique =
    TypeError ('Text "🚫 Movable types can never be Unique!")

Now when someone tries to pass a Movable value directly to a method expecting a Unique they'll get a descriptive compiler error, and (more importantly) if they derive both a Movable and a Unique instance for the same type, GHC will prevent it with a (less descriptive) 'overlapping instances' compiler error.

[aspiwack]

Actually, I think I have an idea. We could use a LinearityToken much like in the Linear constraints proposal.

We could have

data LinearityToken -- abstract
instance Dupable LinearityToken

class Unique a where
  bearsLinearity :: a -> (LinearityToken, a)

Because the LinearityToken is abstract, we can only store it in our structure. But with the Unique witness type class infrastructure, we actually never need to use LinearityToken as a API consumer.

PS: though we can because

instance Unique LinearityToken where
  bearsLinearity = dup2

PPS: You're catch-all instance doesn't work as I assume you intend: being Movable isn't required to select this instance, it will always apply (and return an error) for a type where there isn't a more precise Unique (T …) available. And conversely, it will never apply when there is a more precise, matching, Unique (T …) available, even if T … is Movable.

[aspiwack]

Regarding (1), I like the idea of replacing Ur b by Movable b => … b. This is backward compatible, and easy to implement. Where I don't follow is the motivation for the Immutable type class. What kind of types you'd like to have Immutable which aren't Movable?

Or maybe a more practical way of asking the question is: how do I eliminate an Immutable b (i.e. how do I use the fact that b is immutable in order to make a safe API). With Ur: I seq (this makes sure that there are no unused linear values captured in the closure of the Ur thunk). With Movable, I seq the move-d value. What do I do with an Immutable b?

[Qqwy]

I think I was mistaken in adding a separate Immutable class. I have not been able to come up with examples of types for which yes-Immutable no-Movable makes sense.

[aspiwack]

Ok, then let's go for Movable 🙂 .