From f8db42a08335991d8689d65ce023f43dde40e9b9 Mon Sep 17 00:00:00 2001
From: Tomas Glozar <tglozar@gmail.com>
Date: Mon, 29 May 2023 15:00:43 +0200
Subject: [PATCH] Tests for user function definition support

Includes parser, codegen and semantic analyzer tests.
---
 tests/codegen/llvm/subprog_arguments.ll |  29 ++++++
 tests/codegen/subprog.cpp               |  14 +++
 tests/parser.cpp                        | 115 ++++++++++++++++++++++++
 tests/semantic_analyser.cpp             |  14 +++
 4 files changed, 172 insertions(+)
 create mode 100644 tests/codegen/llvm/subprog_arguments.ll
 create mode 100644 tests/codegen/subprog.cpp

diff --git a/tests/codegen/llvm/subprog_arguments.ll b/tests/codegen/llvm/subprog_arguments.ll
new file mode 100644
index 000000000000..5f242a329907
--- /dev/null
+++ b/tests/codegen/llvm/subprog_arguments.ll
@@ -0,0 +1,29 @@
+; ModuleID = 'bpftrace'
+source_filename = "bpftrace"
+target datalayout = "e-m:e-p:64:64-i64:64-i128:128-n32:64-S128"
+target triple = "bpf-pc-linux"
+
+; Function Attrs: nounwind
+declare i64 @llvm.bpf.pseudo(i64 %0, i64 %1) #0
+
+define internal i64 @add(i64 %0, i64 %1) {
+entry:
+  %"$b" = alloca i64, align 8
+  %"$a" = alloca i64, align 8
+  %2 = bitcast i64* %"$a" to i8*
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* %2)
+  store i64 %0, i64* %"$a", align 8
+  %3 = bitcast i64* %"$b" to i8*
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* %3)
+  store i64 %1, i64* %"$b", align 8
+  %4 = load i64, i64* %"$a", align 8
+  %5 = load i64, i64* %"$b", align 8
+  %6 = add i64 %4, %5
+  ret i64 %6
+}
+
+; Function Attrs: argmemonly nofree nosync nounwind willreturn
+declare void @llvm.lifetime.start.p0i8(i64 immarg %0, i8* nocapture %1) #1
+
+attributes #0 = { nounwind }
+attributes #1 = { argmemonly nofree nosync nounwind willreturn }
diff --git a/tests/codegen/subprog.cpp b/tests/codegen/subprog.cpp
new file mode 100644
index 000000000000..b5c7ad15096f
--- /dev/null
+++ b/tests/codegen/subprog.cpp
@@ -0,0 +1,14 @@
+#include "common.h"
+
+namespace bpftrace {
+namespace test {
+namespace codegen {
+
+TEST(codegen, subprog_arguments)
+{
+  test("fn add($a : int64, $b : int64): int64 { return $a + $b; }", NAME);
+}
+
+} // namespace codegen
+} // namespace test
+} // namespace bpftrace
diff --git a/tests/parser.cpp b/tests/parser.cpp
index 1b632baeed0b..111f56ac9fa3 100644
--- a/tests/parser.cpp
+++ b/tests/parser.cpp
@@ -1971,6 +1971,121 @@ TEST(Parser, invalid_provider)
   test_parse_failure("asdf:xyz:www { }");
 }
 
+TEST(Parser, subprog_only)
+{
+  test("fn f1(): void {}",
+       "Program\n"
+       " f1: void()\n");
+}
+
+TEST(Parser, subprog_probe_mixed)
+{
+  test("i:s:1 {} fn f1(): void {} i:s:1 {} fn f2(): void {}",
+       "Program\n"
+       " f1: void()\n"
+       " f2: void()\n"
+       " interval:s:1\n"
+       " interval:s:1\n");
+}
+
+TEST(Parser, subprog_void_no_args)
+{
+  test("fn f(): void {}",
+       "Program\n"
+       " f: void()\n");
+}
+
+TEST(Parser, subprog_invalid_return_type)
+{
+  test_parse_failure("fn f(): nonexistent {}");
+}
+
+TEST(Parser, subprog_one_arg)
+{
+  test("fn f($a : uint8): void {}",
+       "Program\n"
+       " f: void(unsigned int8 $a)\n");
+}
+
+TEST(Parser, subprog_two_args)
+{
+  test("fn f($a : uint8, $b: uint8): void {}",
+       "Program\n"
+       " f: void(unsigned int8 $a, unsigned int8 $b)\n");
+}
+
+TEST(Parser, subprog_string_arg)
+{
+  test("fn f($a : str_t[16]): void {}",
+       "Program\n"
+       " f: void(string[16] $a)\n");
+}
+
+TEST(Parser, subprog_struct_arg)
+{
+  test("fn f($a: struct x): void {}",
+       "Program\n"
+       " f: void(struct x $a)\n");
+}
+
+TEST(Parser, subprog_union_arg)
+{
+  test("fn f($a : union x): void {}",
+       "Program\n"
+       " f: void(union x $a)\n");
+}
+
+TEST(Parser, subprog_enum_arg)
+{
+  test("fn f($a : enum x): void {}",
+       "Program\n"
+       " f: void(enum x $a)\n");
+}
+
+TEST(Parser, subprog_invalid_arg)
+{
+  test_parse_failure("fn f($x : invalid): void {}");
+}
+
+TEST(Parser, subprog_return)
+{
+  test("fn f(): void { return 1 + 1; }",
+       "Program\n"
+       " f: void()\n"
+       "  return\n"
+       "   +\n"
+       "    int: 1\n"
+       "    int: 1\n");
+}
+
+TEST(Parser, subprog_string)
+{
+  test("fn f(): str_t[16] {}",
+       "Program\n"
+       " f: string[16]()\n");
+}
+
+TEST(Parser, subprog_struct)
+{
+  test("fn f(): struct x {}",
+       "Program\n"
+       " f: struct x()\n");
+}
+
+TEST(Parser, subprog_union)
+{
+  test("fn f(): union x {}",
+       "Program\n"
+       " f: union x()\n");
+}
+
+TEST(Parser, subprog_enum)
+{
+  test("fn f(): enum x {}",
+       "Program\n"
+       " f: enum x()\n");
+}
+
 } // namespace parser
 } // namespace test
 } // namespace bpftrace
diff --git a/tests/semantic_analyser.cpp b/tests/semantic_analyser.cpp
index f5ac8ba604dd..bb4fab562a5f 100644
--- a/tests/semantic_analyser.cpp
+++ b/tests/semantic_analyser.cpp
@@ -2702,6 +2702,20 @@ TEST(semantic_analyser, string_size)
   ASSERT_EQ(var_assign->var->type.GetField(0).type.GetSize(), 6UL);
 }
 
+TEST(semantic_analyser, subprog_return)
+{
+  test("fn f(): void { return; }", 0);
+  test("fn f(): void { return 1; }", 1);
+  test("fn f(): int64 { return; }", 1);
+  test("fn f(): int64 { return 1; }", 0);
+}
+
+TEST(semantic_analyser, subprog_arguments)
+{
+  test("fn f($a : int64): int64 { return $a; }", 0);
+  test("fn f($a : int64): str_t[16] { return $a; }", 1);
+}
+
 class semantic_analyser_btf : public test_btf
 {
 };