Haskell: use last reference to a variable to efficiently create a new variable

Question

This C code can conceptually be described as creating a new array identical to an input array but with 1 as the first element:

int* retire_and_update(int* arr) {
    arr[0] = 1;
    return arr;
}

This is a pure function (wink wink nudge nudge) so long as no further references are made to the input array and its elements. The C type system won't enforce that for us but it seems enforceable in principle.

The code that gcc generates is simple and efficient:

retire_and_update:
    movq    %rdi, %rax
    movl    $1, (%rdi)
    ret

Our function achieves the seemingly impossible by creating a whole new array in constant time and using no additional memory. Nice. Can a Haskell function with array-like input and output be written that could be validly implemented with similar code? Is there a way to express "this is the last reference to this variable" so that a pure function can cannibalize the variable behind the scenes?

If the function gets inlined then nothing interesting needs to happen here so let's suppose the caller and function will be compiled separately.

Answer 1

Though ST monad is not exactly what you describe, in practice you may implement most of that using STUArray. So, a mock up of your code may be something like:

import Control.Monad (forM_)
import Control.Monad.ST (ST)
import Data.Array.Unboxed (UArray)
import Data.Array.ST (STUArray, newArray, readArray, writeArray, runSTUArray)

retire_and_update :: STUArray s Int Int -> ST s (STUArray s Int Int)
retire_and_update arr = do
    writeArray arr 0 1
    return arr

and if you have another function which mutates the array in-place, like:

mutate_inplace :: STUArray s Int Int -> Int -> ST s ()
mutate_inplace arr size = do
    forM_ [2..size - 1] $ \i -> do
        a <- readArray arr (i - 2)
        b <- readArray arr (i - 1)
        writeArray arr i (a + b)

you may bind the two impure function together and call them inside a pure function using runSTUArray:

run :: Int -> UArray Int Int
run size = runSTUArray $ do
    arr <- newArray (0, size - 1) 0
    retire_and_update arr
    mutate_inplace arr size
    return arr

note that run stays pure, and the earlier versions of the returned array are not leaked out anywhere:

\> run 8
array (0,7) [(0,1),(1,0),(2,1),(3,1),(4,2),(5,3),(6,5),(7,8)]