1 files changed, 55 insertions, 49 deletions
diff --git a/lib/tsan/tests/unit/tsan_stack_test.cc b/lib/tsan/tests/unit/tsan_stack_test.cc
index 9aa2967628cf9..92e035d8d0004 100644
--- a/lib/tsan/tests/unit/tsan_stack_test.cc
+++ b/lib/tsan/tests/unit/tsan_stack_test.cc
@@ -17,73 +17,79 @@
 
 namespace __tsan {
 
-static void TestStackTrace(StackTrace *trace) {
-  ThreadState thr(0, 0, 0, 0, 0, 0, 0, 0);
+template <typename StackTraceTy>
+static void TestStackTrace(StackTraceTy *trace) {
+  ThreadState thr(0, 0, 0, 0, 0, 0, 0, 0, 0);
   uptr stack[128];
   thr.shadow_stack = &stack[0];
   thr.shadow_stack_pos = &stack[0];
   thr.shadow_stack_end = &stack[128];
 
-  trace->ObtainCurrent(&thr, 0);
-  EXPECT_EQ(trace->Size(), (uptr)0);
+  ObtainCurrentStack(&thr, 0, trace);
+  EXPECT_EQ(0U, trace->size);
 
-  trace->ObtainCurrent(&thr, 42);
-  EXPECT_EQ(trace->Size(), (uptr)1);
-  EXPECT_EQ(trace->Get(0), (uptr)42);
+  ObtainCurrentStack(&thr, 42, trace);
+  EXPECT_EQ(1U, trace->size);
+  EXPECT_EQ(42U, trace->trace[0]);
 
   *thr.shadow_stack_pos++ = 100;
   *thr.shadow_stack_pos++ = 101;
-  trace->ObtainCurrent(&thr, 0);
-  EXPECT_EQ(trace->Size(), (uptr)2);
-  EXPECT_EQ(trace->Get(0), (uptr)100);
-  EXPECT_EQ(trace->Get(1), (uptr)101);
+  ObtainCurrentStack(&thr, 0, trace);
+  EXPECT_EQ(2U, trace->size);
+  EXPECT_EQ(100U, trace->trace[0]);
+  EXPECT_EQ(101U, trace->trace[1]);
 
-  trace->ObtainCurrent(&thr, 42);
-  EXPECT_EQ(trace->Size(), (uptr)3);
-  EXPECT_EQ(trace->Get(0), (uptr)100);
-  EXPECT_EQ(trace->Get(1), (uptr)101);
-  EXPECT_EQ(trace->Get(2), (uptr)42);
+  ObtainCurrentStack(&thr, 42, trace);
+  EXPECT_EQ(3U, trace->size);
+  EXPECT_EQ(100U, trace->trace[0]);
+  EXPECT_EQ(101U, trace->trace[1]);
+  EXPECT_EQ(42U, trace->trace[2]);
 }
 
-TEST(StackTrace, Basic) {
-  ScopedInRtl in_rtl;
-  StackTrace trace;
-  TestStackTrace(&trace);
-}
+template<typename StackTraceTy>
+static void TestTrim(StackTraceTy *trace) {
+  ThreadState thr(0, 0, 0, 0, 0, 0, 0, 0, 0);
+  const uptr kShadowStackSize = 2 * kStackTraceMax;
+  uptr stack[kShadowStackSize];
+  thr.shadow_stack = &stack[0];
+  thr.shadow_stack_pos = &stack[0];
+  thr.shadow_stack_end = &stack[kShadowStackSize];
 
-TEST(StackTrace, StaticBasic) {
-  ScopedInRtl in_rtl;
-  uptr buf[10];
-  StackTrace trace1(buf, 10);
-  TestStackTrace(&trace1);
-  StackTrace trace2(buf, 3);
-  TestStackTrace(&trace2);
-}
+  for (uptr i = 0; i < kShadowStackSize; ++i)
+    *thr.shadow_stack_pos++ = 100 + i;
 
-TEST(StackTrace, StaticTrim) {
-  ScopedInRtl in_rtl;
-  uptr buf[2];
-  StackTrace trace(buf, 2);
+  ObtainCurrentStack(&thr, 0, trace);
+  EXPECT_EQ(kStackTraceMax, trace->size);
+  for (uptr i = 0; i < kStackTraceMax; i++) {
+    EXPECT_EQ(100 + kStackTraceMax + i, trace->trace[i]);
+  }
 
-  ThreadState thr(0, 0, 0, 0, 0, 0, 0, 0);
-  uptr stack[128];
-  thr.shadow_stack = &stack[0];
-  thr.shadow_stack_pos = &stack[0];
-  thr.shadow_stack_end = &stack[128];
+  ObtainCurrentStack(&thr, 42, trace);
+  EXPECT_EQ(kStackTraceMax, trace->size);
+  for (uptr i = 0; i < kStackTraceMax - 1; i++) {
+    EXPECT_EQ(101 + kStackTraceMax + i, trace->trace[i]);
+  }
+  EXPECT_EQ(42U, trace->trace[kStackTraceMax - 1]);
+}
 
-  *thr.shadow_stack_pos++ = 100;
-  *thr.shadow_stack_pos++ = 101;
-  *thr.shadow_stack_pos++ = 102;
-  trace.ObtainCurrent(&thr, 0);
-  EXPECT_EQ(trace.Size(), (uptr)2);
-  EXPECT_EQ(trace.Get(0), (uptr)101);
-  EXPECT_EQ(trace.Get(1), (uptr)102);
+TEST(StackTrace, BasicVarSize) {
+  VarSizeStackTrace trace;
+  TestStackTrace(&trace);
+}
 
-  trace.ObtainCurrent(&thr, 42);
-  EXPECT_EQ(trace.Size(), (uptr)2);
-  EXPECT_EQ(trace.Get(0), (uptr)102);
-  EXPECT_EQ(trace.Get(1), (uptr)42);
+TEST(StackTrace, BasicBuffered) {
+  BufferedStackTrace trace;
+  TestStackTrace(&trace);
+}
+
+TEST(StackTrace, TrimVarSize) {
+  VarSizeStackTrace trace;
+  TestTrim(&trace);
 }
 
+TEST(StackTrace, TrimBuffered) {
+  BufferedStackTrace trace;
+  TestTrim(&trace);
+}
 
 }  // namespace __tsan