Forward fill first instances of NAs in R data.table

I have a data.table with a column like:

c(58,NA,NA,NA,NA,13,NA,NA,NA,12,23,NA,12)

I would like to fill only the first two NAs following each non-NA value in the column by forwarding the last previous value. The result should be :

c(58,58,58,NA,NA,13,13,13,NA,12,23,23,12)

Any suggestions ?

>Solution :

One crude way in base R:

na_locf_max <- function(x, nmax){
  # Split the vector in sequence of numbers and NAs
  s <- split(x, cumsum(!is.na(x)))

  # Assign the first nmax value to the first value, and trim to match length
  l <- mapply(\(x, y) {
    x[1:nmax+1] <- x[1]
    length(x) <- y
    x
  }, s, lengths(s))

  #Format to vector format
  unlist(l, use.names = FALSE)
}

na_locf_max(x, nmax = 2)
# [1] 58 58 58 NA NA 13 13 13 NA 12 23 23 12