Get the number of elements for Vector and Matrix type

This commit is contained in:
qining 2016-08-23 11:34:24 -04:00
parent b891845b53
commit 7e775aa99e
2 changed files with 19 additions and 0 deletions

View File

@ -162,6 +162,7 @@ class Vector : public Type {
bool IsSame(Type* that) const override;
std::string str() const override;
const Type* element_type() const { return element_type_; }
uint32_t element_count() const { return count_; }
Vector* AsVector() override { return this; }
const Vector* AsVector() const override { return this; }
@ -179,6 +180,7 @@ class Matrix : public Type {
bool IsSame(Type* that) const override;
std::string str() const override;
const Type* element_type() const { return element_type_; }
uint32_t element_count() const { return count_; }
Matrix* AsMatrix() override { return this; }
const Matrix* AsMatrix() const override { return this; }

View File

@ -240,4 +240,21 @@ TEST(Types, FloatWidth) {
}
}
TEST(Types, VectorElementCount) {
auto s32 = std::unique_ptr<Integer>(new Integer(32, true));
for (uint32_t c : {2, 3, 4}) {
auto s32v = std::unique_ptr<Vector>(new Vector(s32.get(), c));
EXPECT_EQ(c, s32v->element_count());
}
}
TEST(Types, MatrixElementCount) {
auto s32 = std::unique_ptr<Integer>(new Integer(32, true));
auto s32v4 = std::unique_ptr<Vector>(new Vector(s32.get(), 4));
for (uint32_t c : {1, 2, 3, 4, 10, 100}) {
auto s32m = std::unique_ptr<Matrix>(new Matrix(s32v4.get(), c));
EXPECT_EQ(c, s32m->element_count());
}
}
} // anonymous namespace