Modified the test code in main.cpp and fixed a segfault