diff --git a/ffi/include/tvm/ffi/container/map.h b/ffi/include/tvm/ffi/container/map.h index fc94ab15a854..7892e2008530 100644 --- a/ffi/include/tvm/ffi/container/map.h +++ b/ffi/include/tvm/ffi/container/map.h @@ -1264,17 +1264,27 @@ inline ObjectPtr MapObj::CreateFromRange(IterType first, IterType last) } uint64_t cap = static_cast(_cap); if (cap < SmallMapObj::kMaxSize) { - return SmallMapObj::CreateFromRange(cap, first, last); - } - uint32_t fib_shift; - uint64_t n_slots; - DenseMapObj::CalcTableSize(cap, &fib_shift, &n_slots); - ObjectPtr obj = DenseMapObj::Empty(fib_shift, n_slots); - for (; first != last; ++first) { - KVType kv(*first); - DenseMapObj::InsertMaybeReHash(std::move(kv), &obj); + if (cap < 2) { + return SmallMapObj::CreateFromRange(cap, first, last); + } + // need to insert to avoid duplicate keys + ObjectPtr obj = SmallMapObj::Empty(cap); + for (; first != last; ++first) { + KVType kv(*first); + SmallMapObj::InsertMaybeReHash(std::move(kv), &obj); + } + return obj; + } else { + uint32_t fib_shift; + uint64_t n_slots; + DenseMapObj::CalcTableSize(cap, &fib_shift, &n_slots); + ObjectPtr obj = DenseMapObj::Empty(fib_shift, n_slots); + for (; first != last; ++first) { + KVType kv(*first); + DenseMapObj::InsertMaybeReHash(std::move(kv), &obj); + } + return obj; } - return obj; } inline void MapObj::InsertMaybeReHash(KVType&& kv, ObjectPtr* map) { diff --git a/ffi/tests/cpp/test_map.cc b/ffi/tests/cpp/test_map.cc index b7c977fd344c..98d8427c23a1 100644 --- a/ffi/tests/cpp/test_map.cc +++ b/ffi/tests/cpp/test_map.cc @@ -356,4 +356,11 @@ TEST(Map, EmptyIter) { // now m0 is dense map with all empty slots EXPECT_EQ(m0.begin(), m0.end()); } + +TEST(Map, DuplicatedKeysInit) { + std::vector> data = {{"a", 1}, {"a", 2}, {"a", 3}}; + Map map(data.begin(), data.end()); + EXPECT_EQ(map.size(), 1); + EXPECT_EQ(map["a"], 3); +} } // namespace