//===- sparse_tensor.c - Test of sparse_tensor APIs -----------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM
// Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

// RUN: mlir-capi-sparse-tensor-test 2>&1 | FileCheck %s

#include "mlir-c/Dialect/SparseTensor.h"
#include "mlir-c/IR.h"
#include "mlir-c/Registration.h"

#include <assert.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

// CHECK-LABEL: testRoundtripEncoding()
static int testRoundtripEncoding(MlirContext ctx) {
  fprintf(stderr, "testRoundtripEncoding()\n");
  // clang-format off
  const char *originalAsm =
    "#sparse_tensor.encoding<{ "
    "dimLevelType = [ \"dense\", \"compressed\", \"singleton\"], "
    "dimOrdering = affine_map<(d0, d1, d2) -> (d0, d1, d2)>, "
    "pointerBitWidth = 32, indexBitWidth = 64 }>";
  // clang-format on
  MlirAttribute originalAttr =
      mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString(originalAsm));
  // CHECK: isa: 1
  fprintf(stderr, "isa: %d\n",
          mlirAttributeIsASparseTensorEncodingAttr(originalAttr));
  MlirAffineMap dimOrdering =
      mlirSparseTensorEncodingAttrGetDimOrdering(originalAttr);
  // CHECK: (d0, d1, d2) -> (d0, d1, d2)
  mlirAffineMapDump(dimOrdering);
  // CHECK: level_type: 0
  // CHECK: level_type: 1
  // CHECK: level_type: 2
  int numLevelTypes = mlirSparseTensorEncodingGetNumDimLevelTypes(originalAttr);
  enum MlirSparseTensorDimLevelType *levelTypes =
      malloc(sizeof(enum MlirSparseTensorDimLevelType) * numLevelTypes);
  for (int i = 0; i < numLevelTypes; ++i) {
    levelTypes[i] =
        mlirSparseTensorEncodingAttrGetDimLevelType(originalAttr, i);
    fprintf(stderr, "level_type: %d\n", levelTypes[i]);
  }
  // CHECK: pointer: 32
  int pointerBitWidth =
      mlirSparseTensorEncodingAttrGetPointerBitWidth(originalAttr);
  fprintf(stderr, "pointer: %d\n", pointerBitWidth);
  // CHECK: index: 64
  int indexBitWidth =
      mlirSparseTensorEncodingAttrGetIndexBitWidth(originalAttr);
  fprintf(stderr, "index: %d\n", indexBitWidth);

  MlirAttribute newAttr = mlirSparseTensorEncodingAttrGet(
      ctx, numLevelTypes, levelTypes, dimOrdering, pointerBitWidth,
      indexBitWidth);
  mlirAttributeDump(newAttr); // For debugging filecheck output.
  // CHECK: equal: 1
  fprintf(stderr, "equal: %d\n", mlirAttributeEqual(originalAttr, newAttr));

  free(levelTypes);
  return 0;
}

int main() {
  MlirContext ctx = mlirContextCreate();
  mlirDialectHandleRegisterDialect(mlirGetDialectHandle__sparse_tensor__(),
                                   ctx);
  if (testRoundtripEncoding(ctx))
    return 1;

  mlirContextDestroy(ctx);
  return 0;
}