Treating return value as an input parameter

Evolution Discussion
#1

I’ve been mulling this over for a while, and I think there is an opportunity to make efficient code as ergonomic as inefficient code, when it comes to reusing existing buffer storage for the result of a calculation.

Suppose you have a type which uses a buffer for storage, and certain operations return that type. For example, a Vector or Matrix with arithmetic operators.

Code like this will of course allocate a new instance with a new buffer for w on each iteration:

for u in vecList {
  for v in vecList {
    let w = u .- v
    // use w
  }
}

Unfortunately so will code like this:

var w = Vector()
for u in vecList {
  for v in vecList {
    w = u .- v
    // use w
  }
}

Even if the result of a calculation is stored to an instance which already exists and has a buffer capable of holding the result, that existing buffer is not used. In this example, the subtraction operator creates a new instance of Vector, which allocates a new buffer to hold the result, regardless of the fact that w already had a perfectly serviceable buffer right there.

In order to reuse the storage buffer in w, we have to write code like this:

var w = Vector()
for u in vecList {
  for v in vecList {
    u.minus(v, storedInto: &w)
    // use w
  }
}

That is unwieldy, and I think we can all agree it is a pattern we do not want to encourage.

If instead it were possible to make the return value of the operator use existing storage when available, then we could keep the natural spelling for the assignment, namely w = u .- v, and still get the efficiency of reusing the same buffer.

Perhaps the declaration of .- might look something like this:

extension Vector where Element: Numeric {
  static func .- (lhs: Self, rhs: Self) -> (result: inout Self = Vector()) {
    let n = lhs.storage.count
    precondition(rhs.storage.count == n)
    
    result.storage.removeAll(keepingCapacity: true)
    result.storage.reserveCapacity(n)
    
    for (x, y) in zip(lhs.storage, rhs.storage) {
      result.storage.append(x - y)
    }
  }
}

The declaration syntax is entirely bikesheddable, but the conceptual idea is that the instance into which the return value will be stored (w), is available as a mutable parameter inside the function body.

The default value for result indicates that if no existing instance is provided (such as in let w = u .- v) then the default is used. If no default were provided, then presumably there would be a separate declaration of the operation with a standard, non-mutable return value for those calls.

Note that this is different from what I understand previous discussions about “inout return values” to mean. I am not tied to any particular spelling or name for the feature I am describing, but I think it is valuable.

1 Like
#2

It sounds like something a compiler should be able to do implicitly. If it can't, extra annotation won't help us much since the compiler will just use the provided allocator anyway. Also, it doesn't get us much further than assigning into local variable/stored property, computed setter is already an easy barrier.

Perhaps we can have one for functions across ABI boundaries, but I'm not sure if we should just use @inlinable.

PS
Another thing is that I'm afraid this pattern will encourage NSZone-style usage. I personally don't like it, but it's not exactly a downside.

1 Like
#3

FWIW: Assuming Vector is a value type with cow semantics, and its storage is a reference type (ie, storage is not another (nested) cow value type like Array), I guess you could do something like the following in current Swift:

infix operator .=: AssignmentPrecedence
infix operator .-: AdditionPrecedence

struct Vector { /* … */ }
struct VectorInstruction { /* … holds references to storages of two vectors and can compute result into a given vector's storage */ }

func .=(lhs: inout Vector, rhs: VectorInstruction) { /* … */ }

func .-(lhs: Vector, rhs: Vector) -> VectorInstruction { /* … */ }

var w = Vector()
for u in vecList {
  for v in vecList {
    w .= u .- v
    // use w
  }
}

Also assuming that the vector-instruction-thing can be inlined and ends up producing no unnecessary overhead.

Horrible proof(?) of concept.

I just threw this together without caring much about anything so it's ugly and bad in all possible ways, but at least it compiles and the output is as I expected:

fileprivate final class VectorStorage {
  var ptr: UnsafeMutableBufferPointer<Float>
  var capacity: Int
  var count: Int
  init(capacity: Int) {
    precondition(capacity > 0)
    self.capacity = capacity
    self.count = 0
    self.ptr = .allocate(capacity: capacity)
    print("init", ptr)
  }
  convenience init<C>(from elements: C, capacity: Int)
    where C: Collection, C.Element == Float
  {
    precondition(capacity >= elements.count)
    self.init(capacity: capacity)
    _ = ptr.initialize(from: elements)
  }
  deinit {
    print("deinit", ptr)
    ptr.deallocate()
  }
}

struct Vector {
  fileprivate var storage: VectorStorage
  var count: Int { storage.count }
  init<C>(from elements: C) where C: Collection, C.Element == Float {
    precondition(!elements.isEmpty)
    self.storage = VectorStorage(from: elements, capacity: elements.count)
  }
  subscript(index: Int) -> Float {
    get {
      precondition((0 ..< count).contains(index))
      return storage.ptr[index]
    }
    set {
      precondition((0 ..< count).contains(index))
      copyStorageIfNeeded(capacity: storage.capacity)
      storage.ptr[index] = newValue
    }
  }
  mutating func append(_ element: Float) {
    if storage.capacity == storage.count {
      storage = VectorStorage(from: storage.ptr, capacity: storage.count*2)
    } else {
      copyStorageIfNeeded(capacity: storage.capacity)
    }
    storage.ptr[storage.count] = element
    storage.count += 1
  }
  mutating func removeAll(keepingCapacity: Bool) {
    if keepingCapacity {
      if !isKnownUniquelyReferenced(&storage) {
        storage = VectorStorage(capacity: storage.capacity)
      } else {
        storage.count = 0
      }
    } else {
      storage = VectorStorage(capacity: 1)
    }
  }
  mutating func reserveCapacity(_ minimumCapacity: Int) {
    if minimumCapacity <= storage.capacity { return }
    copyStorageIfNeeded(capacity: minimumCapacity)
  }
  private mutating func copyStorageIfNeeded(capacity: Int) {
    if !isKnownUniquelyReferenced(&storage) {
      storage = VectorStorage(from: storage.ptr, capacity: capacity)
    }
  }
}
extension Vector {
  static func random(count: Int, in range: Range<Float>) -> Self {
    let elements = (0 ..< count).map { _ in Float.random(in: range) }
    return Self(from: elements)
  }
}


infix operator .=: AssignmentPrecedence
infix operator .-: AdditionPrecedence

struct VectorInstruction {
  enum Operation {
    case add, mul, sub, div
  }
  fileprivate var op: Operation
  fileprivate var a: VectorStorage
  fileprivate var b: VectorStorage
  fileprivate init(op: Operation, _ a: VectorStorage, _ b: VectorStorage) {
    precondition(a.count == b.count)
    self.op = op
    self.a = a
    self.b = b
  }
  func computeResult(into result: inout Vector) {
    result.removeAll(keepingCapacity: true)
    result.reserveCapacity(a.count)
    for (x, y) in zip(a.ptr, b.ptr) {
      result.append(x - y) // Could switch on op cases here.
    }
  }
}

func .=(lhs: inout Vector, rhs: VectorInstruction) {
  rhs.computeResult(into: &lhs)
}

func .-(lhs: Vector, rhs: Vector) -> VectorInstruction {
  return .init(op: .sub, lhs.storage, rhs.storage)
}

func test() {
  let vecList = (0 ..< 7).map { _ in Vector.random(count: 3, in: -3 ..< 3) }
  var w = Vector.random(count: 3, in: -3 ..< 3)
  for u in vecList {
    for v in vecList {
      w .= u .- v
      print("--")
      print(u.storage.ptr.map({ String($0) }).joined(separator: ", "))
      print(v.storage.ptr.map({ String($0) }).joined(separator: ", "))
      print(" ", w.storage.ptr.map({ String($0) }).joined(separator: ", "))
    }
  }
}
test()

The VectorInstruction-thing can be made composable, to handle expressions of more than two terms, etc.