Some functionality improvements