String.sk

/**
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

module Unsafe;

@cpp_runtime("SKIP_Unsafe_string_utf8_size")
native fun string_utf8_size(s: String): Int;

// This is only parameterized on T because changing that to UInt8 causes a
// compiler error: T25454827
@cpp_runtime("SKIP_Unsafe_string_utf8_get")
native fun string_utf8_get<T>(s: String, i: Int): T;

module end;

module String;

extension class .String {
  fun utf8(): IndexedSequence<UInt8> {
    Utf8View<UInt8>(this)
  }
}

// This is only parameterized on T because changing that to UInt8 causes a
// compiler error: T25454827
private class Utf8View<T>(s: String) extends IndexedSequence<T> {
  fun size(): Int {
    Unsafe.string_utf8_size(this.s)
  }

  private fun unsafe_get(i: Int): T {
    Unsafe.string_utf8_get(this.s, i)
  }

  // Needed by IndexedSequence<>
  fun chill(): this {
    this
  }

  fun get(index: Int): T {
    if (index.uge(this.size())) throwOutOfBounds();
    this.unsafe_get(index);
  }

  fun values(): mutable Iterator<T> {
    for (i in Range(0, this.size())) {
      yield this.unsafe_get(i);
    }
  }
}

private fun blockIterator<T>(
  input: Sequence<T>,
  blockSize: Int,
): mutable Iterator<Array<T>> {
  loop {
    // peel off a block from the input
    pred = input.take(blockSize).collect(Array);
    if (pred.isEmpty()) yield break;
    !input = input.drop(blockSize);
    yield pred;
  }
}

// TODO: need an annotation which says "don't reorder these fields"
private mutable class ConvertBuffer{
  mutable input: readonly Array<UInt8>,
  output: mutable Array<UInt8> = Array::mfill(8192, UInt8::truncate(0)),
  pivot: mutable Array<UInt8> = Array::mfill(8192, UInt8::truncate(0)),
  mutable sourceEncoder: Runtime.NonGCPointer = Runtime.NonGCPointer::zero(),
  mutable targetEncoder: Runtime.NonGCPointer = Runtime.NonGCPointer::zero(),
  mutable inputUsed: Int = 0,
  mutable outputUsed: Int = 0,
  mutable pivotSource: Int = 0,
  mutable pivotTarget: Int = 0,
} {
  @cpp_runtime
  mutable native fun setup(
    inputEncoding: String,
    outputEncoding: String,
    throwOnBadCharacter: Bool,
  ): void;

  @cpp_runtime
  mutable native fun teardown(): void;

  @cpp_runtime
  mutable native fun convert(): void;

  @cpp_runtime
  mutable native fun flush(): void;

  readonly fun empty(): Bool {
    this.inputUsed == this.input.size()
  }
}

fun convert(
  input_: Sequence<UInt8>,
  inputEncoding: Encoding,
  outputEncoding: Encoding,
  throwOnBadCharacter: Bool = false,
): mutable Iterator<UInt8> {
  input = blockIterator(input_, 4096);
  inputData = {
    input.next() match {
    | None() -> yield break
    | Some(tmp) -> tmp
    }
  };

  converter = mutable ConvertBuffer{input => inputData};
  converter.setup(
    inputEncoding.code(),
    outputEncoding.code(),
    throwOnBadCharacter,
  );
  try {
    loop {
      // Yield all the output
      for (i in Range(0, converter.outputUsed)) {
        yield converter.output[i];
      };

      converter.convert();

      // See if we need more input
      if (converter.empty()) {
        converter.!input = {
          input.next() match {
          | None() -> break void
          | Some(tmp) ->
            converter.!inputUsed = 0;
            tmp
          }
        }
      };
    };

    // Need to flush the remaining output
    loop {
      for (i in Range(0, converter.outputUsed)) {
        yield converter.output[i];
      };

      converter.flush();
      if (converter.outputUsed == 0) break void;
      for (i in Range(0, converter.outputUsed)) {
        yield converter.output[i];
      }
    };
  } catch {
  | e @ _ ->
    converter.teardown();
    throw e
  };

  converter.teardown();
}