Repository: pysmt/pysmt
Pull request #375 (2016-10-19 19:46:26)

Handling of Sorts
Introduce the support for custom sorts (PySMTTypes).

Two new classes are introduced: _Type and PartialType
- PartialType is basically a partial function, that is used to represent
  the concept of define-sort. Binding to the actual sort is post-poned,
  by using a lambda expression, internally.
- _TypeDecl: This class represents a Type declaration, and as such
  cannot be used directly to instantiate a Symbol. This capture the
  semantics of declare-sort. A wrapper Type() is given to simplify its
  use, and making 0-arity sorts a special case.

Examples:

The following three statements are equivalent:
- Type("Colors")
- Type("Colors", 0)
- _TypeDecl("Colors", 0)()

In particular, the last case shows that instantiating a TypeDeclaration
returns a PySMTType:

  Type(Pair, 2)(Int, Int)

Type declarations and Type instances are memoized.

This commit also simplifies the handling of base types (Bool, Int,
Real), by making them **instances** of PySMTType (and not subclasses).

Other changes:
- formula.py: typename for Symbols must by an instance of PySMTTypes
- type_checker.py: Simplified handling of Equals. This also simplifies
  the handling of custom types
- environment.py: Introduce the TypeManager object
- PySMTType.type_id was removed. 
- Types are now accessible from pysmt.shortcuts

marcogario: @mikand , @agriggio : Is this API for sorts reasonable? My understanding is that declare-sort is more important than define-sort. In practice, we rarely need a define-sort, since we can build the more complex type directly. 
marcogario: This implementation is not thread-safe anymore (due to the use of auto-generated ids for comparison). This can be probably fixed, by changing the equality check.

The use of custom sorts probably requires moving (part of) the typing information within the environment, and extending contextualization methods accordingly.

Patch:
diff --git a/pysmt/environment.py b/pysmt/environment.py
index e7a19ee..9e5b66d 100644
--- a/pysmt/environment.py
+++ b/pysmt/environment.py
@@ -29,6 +29,7 @@ import pysmt.oracles
 import pysmt.formula
 import pysmt.factory
 import pysmt.decorators
+import pysmt.typing
 
 
 class Environment(object):
@@ -42,6 +43,7 @@ class Environment(object):
     """
     TypeCheckerClass = pysmt.type_checker.SimpleTypeChecker
     FormulaManagerClass = pysmt.formula.FormulaManager
+    TypeManagerClass = pysmt.typing.TypeManager
     SimplifierClass = pysmt.simplifier.Simplifier
     #SubstituterClass = pysmt.substituter.MSSubstituter
     SubstituterClass = pysmt.substituter.MGSubstituter
@@ -52,6 +54,7 @@ class Environment(object):
     SizeOracleClass = pysmt.oracles.SizeOracle
     AtomsOracleClass = pysmt.oracles.AtomsOracle
 
+
     def __init__(self):
         self._stc = self.TypeCheckerClass(self)
         self._formula_manager = self.FormulaManagerClass(self)
@@ -65,6 +68,7 @@ class Environment(object):
         self._fvo = self.FreeVarsOracleClass(self)
         self._sizeo = self.SizeOracleClass(self)
         self._ao = self.AtomsOracleClass(self)
+        self._type_manager = self.TypeManagerClass(self)
 
         self._factory = None
         # Configurations
@@ -78,6 +82,10 @@ class Environment(object):
     def formula_manager(self):
         return self._formula_manager
 
+    @property
+    def type_manager(self):
+        return self._type_manager
+
     @property
     def simplifier(self):
         return self._simplifier
diff --git a/pysmt/formula.py b/pysmt/formula.py
index faca388..4f9e6d3 100644
--- a/pysmt/formula.py
+++ b/pysmt/formula.py
@@ -90,6 +90,8 @@ class FormulaManager(object):
     def _create_symbol(self, name, typename=types.BOOL):
         if len(name) == 0:
             raise PysmtValueError("Empty string is not a valid name")
+        if not isinstance(typename, types.PySMTType):
+            raise PysmtValueError("typename must be a PySMTType.")
         n = self.create_node(node_type=op.SYMBOL,
                              args=tuple(),
                              payload=(name, typename))
@@ -939,9 +941,7 @@ class FormulaManager(object):
               obtain f_b that is the formula f_a expressed on the
               FormulaManager b : f_b = b.normalize(f_a)
         """
-        # TODO: name clash with formula normalization
-        # TODO: Move this out of the manager and into ad-hoc function
-        normalizer = IdentityDagWalker(self.env)
+        normalizer = FormulaContextualizer(self.env)
         return normalizer.walk(formula)
 
     def _polymorph_args_to_tuple(self, args):
@@ -969,3 +969,23 @@ class FormulaManager(object):
             return False
 
 #EOC FormulaManager
+
+
+class FormulaContextualizer(IdentityDagWalker):
+    """Helper class to recreate a formula within a new environment."""
+
+    def __init__(self, env=None):
+        IdentityDagWalker.__init__(self, env=env)
+        self.type_normalize = self.env.type_manager.normalize
+
+    def walk_symbol(self, formula, **kwargs):
+        # Recreate the Symbol taking into account the type information
+        ty = formula.symbol_type()
+        newty = self.type_normalize(ty)
+        return self.mgr.Symbol(formula.symbol_name(), newty)
+
+    def walk_array_value(self, formula, args, **kwargs):
+        # Recreate the ArrayValue taking into account the type information
+        assign = dict(zip(args[1::2], args[2::2]))
+        ty = self.type_normalize(formula.array_value_index_type())
+        return self.mgr.Array(ty, args[0], assign)
diff --git a/pysmt/shortcuts.py b/pysmt/shortcuts.py
index 940b7ca..0c2b219 100644
--- a/pysmt/shortcuts.py
+++ b/pysmt/shortcuts.py
@@ -32,12 +32,16 @@ environment is used (this is the default behavior of
 import warnings
 warnings.simplefilter('default')
 
-import pysmt.typing as types
 import pysmt.configuration as config
 import pysmt.environment
+import pysmt.typing as types
 import pysmt.smtlib.parser
 import pysmt.smtlib.script
 
+# Import types from shortcuts
+from pysmt.typing import INT, BOOL, REAL, BVType, FunctionType, ArrayType, Type
+assert INT or BOOL or REAL or BVType or FunctionType or ArrayType or Type
+
 
 def get_env():
     """Returns the global environment.
diff --git a/pysmt/smtlib/parser.py b/pysmt/smtlib/parser.py
index 8afc151..402cf20 100644
--- a/pysmt/smtlib/parser.py
+++ b/pysmt/smtlib/parser.py
@@ -1112,7 +1112,7 @@ class SmtLibParser(object):
         return SmtLibCommand(current, [name, [], rtype])
 
     def _cmd_get_assertions(self, current, tokens):
-        """(get_assertions)"""
+        """(get-assertions)"""
         self.parse_atoms(tokens, current, 0)
         return SmtLibCommand(current, [])
 
diff --git a/pysmt/type_checker.py b/pysmt/type_checker.py
index 6c3863d..8754e24 100644
--- a/pysmt/type_checker.py
+++ b/pysmt/type_checker.py
@@ -175,15 +175,9 @@ class SimpleTypeChecker(walkers.DagWalker):
 
     def walk_math_relation(self, formula, args, **kwargs):
         #pylint: disable=unused-argument
-        if args[0].is_real_type():
-            return self.walk_type_to_type(formula, args, REAL, BOOL)
-        if args[0].is_int_type():
-            return self.walk_type_to_type(formula, args, INT, BOOL)
         if args[0].is_bv_type():
             return self.walk_bv_to_bool(formula, args)
-        if args[0].is_array_type():
-            return self.walk_type_to_type(formula, args, args[0], BOOL)
-        return None
+        return self.walk_type_to_type(formula, args, args[0], BOOL)
 
     def walk_ite(self, formula, args, **kwargs):
         assert formula is not None
diff --git a/pysmt/typing.py b/pysmt/typing.py
index 4c897b3..10c5447 100644
--- a/pysmt/typing.py
+++ b/pysmt/typing.py
@@ -31,117 +31,146 @@ on a factory service. Each BitVector width is represented by a
 different instance of BVType.
 
 """
+from pysmt.exceptions import PysmtValueError
 
-# Global dictionary of types, used to store the singletons
-__CUSTOM_TYPES__ = {}
-__BV_TYPES__ = {}
-__ARRAY_TYPES__ = {}
 
 class PySMTType(object):
-    """Abstract class for representing a type within pySMT."""
+    """Class for representing a type within pySMT.
 
-    def __init__(self, type_id=-1):
-        self.type_id = type_id
+    Instances of this class are used to represent sorts.
+    The subclass FunctionType is used to represent function declarations.
+
+    """
+
+    def __init__(self, basename=None, args=None):
+        self.basename = basename
+        self.args = args
+        self.arity = len(args) if args else 0
+        if self.args:
+            args = "{%s}" % ", ".join(str(a) for a in self.args)
+        else:
+            args = ""
+        if basename:
+            self.name = basename + args
+        else:
+            self.name = None
 
     def is_bool_type(self):
         return False
 
-    def is_int_type(self):
+    def is_real_type(self):
         return False
 
-    def is_real_type(self):
+    def is_int_type(self):
         return False
 
     def is_bv_type(self, width=None):
         #pylint: disable=unused-argument
         return False
 
-    def is_function_type(self):
+    def is_array_type(self):
         return False
 
-    def is_array_type(self):
+    def is_function_type(self):
         return False
 
     def __hash__(self):
-        return self.type_id
+        return hash(self.name)
 
     def __eq__(self, other):
         if other is None:
             return False
-        return self.type_id == other.type_id
+        if self is other:
+            return True
+        if self.basename == other.basename:
+            return self.args == other.args
+        return False
 
     def __ne__(self, other):
         return not self.__eq__(other)
 
     def __repr__(self):
-        return self.__class__.__name__
+        if self.name is None:
+            return self.__class__.__name__
+        return self.name
+
+    def as_smtlib(self, funstyle=True):
+        name = self.name
+        if self.args:
+            args = " ".join([arg.as_smtlib(funstyle=False) \
+                             for arg in self.args])
+            name = "(" + self.basename + " " + args + ")"
+        if funstyle:
+            return "() %s" % name
+        else:
+            return name
 
+    def __str__(self):
+        return self.name
+
+# EOC PySMTType
 
-class BooleanType(PySMTType):
+# Basic Types Declarations
+class _BoolType(PySMTType):
     def __init__(self):
-        PySMTType.__init__(self, type_id = 0)
+        PySMTType.__init__(self, basename="Bool", args=None)
 
     def is_bool_type(self):
         return True
 
-    def as_smtlib(self, funstyle=True):
-        if funstyle:
-            return "() Bool"
-        else:
-            return "Bool"
-
-    def __str__(self):
-        return "Bool"
+class _IntType(PySMTType):
+    def __init__(self):
+        PySMTType.__init__(self, basename="Int", args=None)
 
+    def is_int_type(self):
+        return True
 
-class RealType(PySMTType):
+class _RealType(PySMTType):
     def __init__(self):
-        PySMTType.__init__(self, type_id = 1)
+        PySMTType.__init__(self, basename="Real", args=None)
 
     def is_real_type(self):
         return True
 
-    def as_smtlib(self, funstyle=True):
-        if funstyle:
-            return "() Real"
-        else:
-            return "Real"
+# End Basic Types Declarations
 
-    def __str__(self):
-        return "Real"
 
+class _ArrayType(PySMTType):
+    """Internal class used to represent an Array type.
 
-class IntType(PySMTType):
-    def __init__(self):
-        PySMTType.__init__(self, type_id = 2)
+    This class should not be instantiated directly, but the factory
+    method ArrayType should be used instead.
+    """
 
-    def is_int_type(self):
-        return True
+    _instances = {}
 
-    def as_smtlib(self, funstyle=True):
-        if funstyle:
-            return "() Int"
-        else:
-            return "Int"
+    def __init__(self, index_type, elem_type):
+        PySMTType.__init__(self,
+                           basename="Array",
+                           args=(index_type, elem_type))
 
-    def __str__(self):
-        return "Int"
+    @property
+    def elem_type(self):
+        """Returns the element type.
 
+        E.g.,  A: (Array Int Real)
+        Returns RealType.
+        """
+        return self.args[1]
 
-def BVType(width=32):
-    """Returns the singleton associated to the BV type for the given width.
+    @property
+    def index_type(self):
+        """Returns the index type.
 
-    This function takes care of building and registering the type
-    whenever needed. To see the functions provided by the type look at
-    _BVType.
-    """
-    key = width
-    if key in __BV_TYPES__:
-        return  __BV_TYPES__[key]
+        E.g.,  A: (Array Int Real)
+        Returns IntType.
+        """
+        return self.args[0]
+
+    def is_array_type(self):
+        return True
 
-    res = _BVType(width=width)
-    __BV_TYPES__[key] = res
-    return res
+# EOC _ArrayType
 
 
 class _BVType(PySMTType):
@@ -150,8 +179,11 @@ class _BVType(PySMTType):
     This class should not be instantiated directly, but the factory
     method BVType should be used instead.
     """
+
+    _instances = {}
+
     def __init__(self, width=32):
-        PySMTType.__init__(self, type_id = 3)
+        PySMTType.__init__(self, basename="BV{%d}" % width, args=None)
         self._width = width
 
     @property
@@ -169,48 +201,17 @@ class _BVType(PySMTType):
         else:
             return "(_ BitVec %d)" % self.width
 
-    def __str__(self):
-        return "BV{%d}" % self.width
-
     def __eq__(self, other):
-        if other is None:
-            return False
-        if self.type_id != other.type_id:
-            return False
-        if self.width != other.width:
-            return False
-        return True
-
-    def __ne__(self, other):
-        return not self.__eq__(other)
+        if PySMTType.__eq__(self, other):
+            return True
+        if other is not None and other.is_bv_type():
+            return self.width == other.width
+        return False
 
     def __hash__(self):
-        return hash(self.type_id + self.width)
-
-
-# FunctionType is a Factory that returns a _FunctionType
-def FunctionType(return_type, param_types):
-    """Returns the singleton of the Function type with the given arguments.
-
-    This function takes care of building and registering the type
-    whenever needed. To see the functions provided by the type look at
-    _FunctionType
+        return hash(self.width)
 
-    Note: If the list of parameters is empty, the function is
-    equivalent to the return type.
-    """
-    param_types = tuple(param_types)
-    key = (return_type, param_types)
-    # 0-arity functions types are equivalent to the return type
-    if len(param_types) == 0:
-        return return_type
-    if key in __CUSTOM_TYPES__:
-        return  __CUSTOM_TYPES__[key]
-
-    res = _FunctionType(return_type=return_type,
-                        param_types=param_types)
-    __CUSTOM_TYPES__[key] = res
-    return res
+# EOC _BVType
 
 
 class _FunctionType(PySMTType):
@@ -219,11 +220,30 @@ class _FunctionType(PySMTType):
     This class should not be instantiated directly, but the factory
     method FunctionType should be used instead.
     """
+
+    _instances = {}
+
     def __init__(self, return_type, param_types):
-        PySMTType.__init__(self, type_id = 4)
+        PySMTType.__init__(self)
         self._return_type = return_type
-        self._param_types = param_types
+        self._param_types = tuple(param_types)
         self._hash = hash(return_type) + sum(hash(p) for p in param_types)
+        # Note:
+
+        # An underlying assumption of this module is that
+        # PySMTType.args can be used as key to identify a given type
+        # instance. This means that all subtypes are accessible
+        # through args (similarly as how we do FNode.args).
+        #
+        # This means that
+        #  - Hashing can use args as a key
+        #  - Navigating the type tree (e.g., during normalization)
+        #    only works on args.
+        #
+        # In order to make this possible, we need to combine the
+        # return typ and param_types for FunctionType.
+        self.args = (self._return_type,) + self.param_types
+        self.arity = len(self.args)
         return
 
     @property
@@ -265,108 +285,237 @@ class _FunctionType(PySMTType):
     def __eq__(self, other):
         if other is None:
             return False
-        if self.type_id != other.type_id:
-            return False
-        if id(self) == id(other):
+        if self is other:
             return True
-        if self.return_type != other.return_type:
-            return False
-        return self.param_types == other.param_types
-
-    def __ne__(self, other):
-        return not self.__eq__(other)
+        if other.is_function_type():
+            if self.return_type == other.return_type and\
+               self.param_types == other.param_types:
+                return True
+        return False
 
     def __hash__(self):
         return self._hash
 
 
-# ArrayType is a Factory that returns a _ArrayType
-def ArrayType(index_type, elem_type):
-    """Returns the singleton of the Array type with the given arguments.
+class _TypeDecl(object):
+    """Create a new Type Declaration (sort).
 
-    This function takes care of building and registering the type
-    whenever needed. To see the functions provided by the type look at
-    _ArrayType
+    This is equivalent to the SMT-LIB command "declare-sort".
+    NOTE: This object is **not** a Type, but a Type Declaration.
     """
-    key = (index_type, elem_type)
-    if key in __ARRAY_TYPES__:
-        return  __ARRAY_TYPES__[key]
 
-    res = _ArrayType(index_type, elem_type)
-    __ARRAY_TYPES__[key] = res
-    return res
+    def __init__(self, type_manager, name, arity):
+        self.typemgr = type_manager
+        self.name = name
+        self.arity = arity
 
+    def __call__(self, *args):
+        return self.typemgr.get_type_instance(self, *args)
 
-class _ArrayType(PySMTType):
-    """Internal class used to represent an Array type.
+# EOC _TypeDecl
 
-    This class should not be instantiated directly, but the factory
-    method ArrayType should be used instead.
+
+class PartialType(object):
+    """PartialType allows to delay the definition of a Type.
+
+    A partial type is equivalent to SMT-LIB "define-sort" command.
     """
-    def __init__(self, index_type, elem_type):
-        PySMTType.__init__(self, type_id = 5)
-        self._index_type = index_type
-        self._elem_type = elem_type
-        self._hash = hash(index_type) + hash(elem_type)
-        return
+    def __init__(self, name, definition):
+        self.name = name
+        self.definition = definition
 
-    @property
-    def elem_type(self):
-        """Returns the element type.
+    def __str__(self):
+        return "PartialType(%s)" % (self.name)
 
-        E.g.,  A: (Array Int Real)
-        Returns RealType.
-        """
-        return self._elem_type
+    def __call__(self, *args):
+        return self.definition(*args)
 
-    @property
-    def index_type(self):
-        """Returns the index type.
+#
+# Constructors
+#
 
-        E.g.,  A: (Array Int Real)
-        Returns IntType.
-        """
-        return self._index_type
 
-    def as_smtlib(self, funstyle=True):
-        itype = self.index_type.as_smtlib(False)
-        etype = self.elem_type.as_smtlib(False)
 
-        if funstyle:
-            return "() (Array %s %s)" % (itype, etype)
-        else:
-            return "(Array %s %s)" % (itype, etype)
+#
+# Singletons for the basic types
+#
+BOOL = _BoolType()
+REAL = _RealType()
+INT =  _IntType()
+PYSMT_TYPES = frozenset([BOOL, REAL, INT])
 
-    def __str__(self):
-        return "Array{%s, %s}" % (self.index_type, self.elem_type)
+# Helper Constants
+BV1, BV8, BV16, BV32, BV64, BV128 = [_BVType(i) for i in [1, 8, 16, 32, 64, 128]]
+ARRAY_INT_INT = _ArrayType(INT,INT)
 
-    def is_array_type(self):
-        return True
 
-    def __eq__(self, other):
-        if other is None:
-            return False
-        if self.type_id != other.type_id:
-            return False
-        if id(self) == id(other):
-            return True
-        if self.index_type != other.index_type:
-            return False
-        return self.elem_type == other.elem_type
+class TypeManager(object):
 
-    def __ne__(self, other):
-        return not self.__eq__(other)
+    def __init__(self, environment):
+        self._bv_types = {}
+        self._function_types = {}
+        self._array_types = {}
+        self._custom_types = {}
+        self._custom_types_decl = {}
+        self.load_global_types()
+        self.environment = environment
 
-    def __hash__(self):
-        return self._hash
+    def load_global_types(self):
+        """Register basic global types within the TypeManager."""
+        for bvtype in (BV1, BV8, BV16, BV32, BV64, BV128):
+            self._bv_types[bvtype.width] = bvtype
+        self._array_types[(INT, INT)] = ARRAY_INT_INT
 
+    def BVType(self, width=32):
+        """Returns the singleton associated to the BV type for the given width.
 
-# Singletons for the basic types
-BOOL = BooleanType()
-REAL = RealType()
-INT = IntType()
+        This function takes care of building and registering the type
+        whenever needed. To see the functions provided by the type look at
+        _BVType.
+        """
+        try:
+            ty = self._bv_types[width]
+        except KeyError:
+            ty = _BVType(width=width)
+            self._bv_types[width] = ty
+        return ty
+
+    def FunctionType(self, return_type, param_types):
+        """Returns the singleton of the Function type with the given arguments.
+
+        This function takes care of building and registering the type
+        whenever needed. To see the functions provided by the type look at
+        _FunctionType
+
+        Note: If the list of parameters is empty, the function is
+        equivalent to the return type.
+        """
+        param_types = tuple(param_types)
+        key = (return_type, param_types)
+        # 0-arity function types are equivalent to the return type
+        if len(param_types) == 0:
+            return return_type
+        try:
+            ty = self._function_types[key]
+        except KeyError:
+            ty = _FunctionType(return_type=return_type,
+                               param_types=param_types)
+            self._function_types[key] = ty
+        return ty
+
+    def ArrayType(self, index_type, elem_type):
+        """Returns the singleton of the Array type with the given arguments.
+
+        This function takes care of building and registering the type
+        whenever needed. To see the functions provided by the type look at
+        _ArrayType
+        """
+        key = (index_type, elem_type)
+        try:
+            ty = self._array_types[key]
+        except KeyError:
+            ty = _ArrayType(index_type, elem_type)
+            self._array_types[key] = ty
+        return ty
+
+    def Type(self, name, arity=0):
+        """Creates a new Type Declaration (sort declaration).
+
+        This is equivalent to the SMT-LIB command "declare-sort".  For
+        sorts of arity 0, we return a PySMTType, all other sorts need to
+        be instantiated.
+
+        See class _Type.
+        """
 
-# Helper Constants
-PYSMT_TYPES = frozenset([BOOL, REAL, INT])
-BV1, BV8, BV16, BV32, BV64, BV128 = [BVType(i) for i in [1, 8, 16, 32, 64, 128]]
-ARRAY_INT_INT = ArrayType(INT,INT)
+        try:
+            td = self._custom_types_decl[name]
+            if td.arity != arity:
+                raise PysmtValueError("Type %s previously declared with arity "\
+                                      " %d." % (name, td.arity))
+        except KeyError:
+            td = _TypeDecl(self, name, arity)
+            self._custom_types_decl[name] = td
+
+        if td.arity == 0:
+            # Automatically instantiate 0-arity types
+            return self.get_type_instance(td)
+        return td
+
+    def get_type_instance(self, type_decl, *args):
+        if not all(isinstance(t, PySMTType) for t in args):
+            raise PysmtValueError("Trying to instantiate %s with non-type args."\
+                                  % str(type_decl))
+        key = (type_decl, tuple(args)) if args is not None else type_decl
+        try:
+            ty = self._custom_types[key]
+        except KeyError:
+            ty = PySMTType(basename=type_decl.name, args=args)
+            self._custom_types[key] = ty
+        return ty
+
+    def normalize(self, type_):
+        """Recursively recreate the given type within the manager.
+
+        This proceeds iteratively on the structure of the type tree.
+        """
+        stack = [type_]
+        typemap = {}
+        while stack:
+            ty = stack.pop()
+            if ty.arity == 0:
+                if ty.is_bool_type() or ty.is_int_type() or ty.is_real_type():
+                    myty = ty
+                elif ty.is_bv_type():
+                    myty = self.BVType(ty.width)
+                else:
+                    myty = self.Type(ty.basename, arity=0)
+                typemap[ty] = myty
+            else:
+                missing = [subtype for subtype in ty.args\
+                           if subtype not in typemap]
+                if missing:
+                    # At least one type still needs to be converted
+                    stack.append(ty)
+                    stack += missing
+                else:
+                    if ty.is_array_type():
+                        index_type = typemap[ty.index_type]
+                        elem_type = typemap[ty.elem_type]
+                        myty = self.ArrayType(index_type, elem_type)
+                    elif ty.is_function_type():
+                        param_types = (typemap[a] for a in ty.param_types)
+                        return_type = typemap[ty.return_type]
+                        myty = self.FunctionType(return_type, param_types)
+                    else:
+                        # Custom Type
+                        typedecl = self.Type(type_.basename, type_.arity)
+                        new_args = (typemap[a] for a in type_.args)
+                        myty = self.get_type_instance(typedecl, new_args)
+                    typemap[ty] = myty
+        return typemap[type_]
+
+# EOC TypeManager
+
+import pysmt.environment
+
+
+def BVType(width=32):
+    """Returns the BV type for the given width."""
+    mgr = pysmt.environment.get_env().type_manager
+    return mgr.BVType(width=width)
+
+def FunctionType(return_type, param_types):
+    """Returns Function Type with the given arguments."""
+    mgr = pysmt.environment.get_env().type_manager
+    return mgr.FunctionType(return_type=return_type, param_types=param_types)
+
+def ArrayType(index_type, elem_type):
+    """Returns the Array type with the given arguments."""
+    mgr = pysmt.environment.get_env().type_manager
+    return mgr.ArrayType(index_type=index_type, elem_type=elem_type)
+
+def Type(name, arity=0):
+    """Returns the Type Declaration with the given name (sort declaration)."""
+    mgr = pysmt.environment.get_env().type_manager
+    return mgr.Type(name=name, arity=arity)
diff --git a/pysmt/walkers/identitydag.py b/pysmt/walkers/identitydag.py
index ca84bcb..e2cdcef 100644
--- a/pysmt/walkers/identitydag.py
+++ b/pysmt/walkers/identitydag.py
@@ -95,8 +95,7 @@ class IdentityDagWalker(DagWalker):
     def walk_function(self, formula, args, **kwargs):
         # We re-create the symbol name
         old_name = formula.function_name()
-        new_name = self.mgr.Symbol(old_name.symbol_name(),
-                                   old_name.symbol_type())
+        new_name = self.walk_symbol(old_name)
         return self.mgr.Function(new_name, args)
 
     def walk_toreal(self, formula, args, **kwargs):


Test patch:
diff --git a/pysmt/test/test_formula.py b/pysmt/test/test_formula.py
index a3f0413..7b6caad 100644
--- a/pysmt/test/test_formula.py
+++ b/pysmt/test/test_formula.py
@@ -1015,19 +1015,6 @@ class TestFormulaManager(TestCase):
         self.assertTrue(and_x_x in self.mgr)
         self.assertFalse(and_y_y in self.mgr)
 
-    def test_typing(self):
-        self.assertTrue(BOOL.is_bool_type())
-        self.assertFalse(BOOL.is_function_type())
-
-        self.assertTrue(REAL.is_real_type())
-        self.assertFalse(REAL.is_bool_type())
-
-        self.assertTrue(INT.is_int_type())
-        self.assertFalse(INT.is_real_type())
-
-        self.assertTrue(self.ftype.is_function_type())
-        self.assertFalse(self.ftype.is_int_type())
-
     def test_array_value(self):
         a1 = self.mgr.Array(INT, self.mgr.Int(0))
         a2 = self.mgr.Array(INT, self.mgr.Int(0),
diff --git a/pysmt/test/test_sorts.py b/pysmt/test/test_sorts.py
new file mode 100644
index 0000000..f08ab04
--- /dev/null
+++ b/pysmt/test/test_sorts.py
@@ -0,0 +1,176 @@
+#
+# This file is part of pySMT.
+#
+#   Copyright 2014 Andrea Micheli and Marco Gario
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
+from six import StringIO
+
+from pysmt.shortcuts import FreshSymbol, EqualsOrIff, Select, TRUE, FALSE, Function
+from pysmt.shortcuts import Array, BV
+from pysmt.typing import INT, BOOL, REAL
+from pysmt.typing import Type, ArrayType, FunctionType, BVType
+from pysmt.typing import PartialType
+from pysmt.test import TestCase, main
+from pysmt.smtlib.parser import SmtLibParser
+from pysmt.exceptions import PysmtValueError, PysmtTypeError
+
+
+SMTLIB_SRC="""\
+(define-sort Set (T) (Array T Bool))
+(define-sort I () Int)
+
+(declare-const s1 (Set I))
+(declare-const a I)
+(declare-const b Int)
+
+(assert (= (select s1 a) true))
+(assert (= (select s1 b) false))
+(check-sat)
+"""
+
+class TestSorts(TestCase):
+
+    def test_smtlib_sort(self):
+        # parser = SmtLibParser()
+        # buf = StringIO(SMTLIB_SRC)
+        # script = parser.get_script(buf)
+        pass
+
+    def test_basic_types(self):
+        self.assertTrue(BOOL.is_bool_type())
+        self.assertFalse(BOOL.is_function_type())
+
+        self.assertTrue(REAL.is_real_type())
+        self.assertFalse(REAL.is_bool_type())
+
+        self.assertTrue(INT.is_int_type())
+        self.assertFalse(INT.is_real_type())
+
+        ftype = FunctionType(REAL, [REAL, REAL])
+        self.assertTrue(ftype.is_function_type())
+        self.assertFalse(ftype.is_int_type())
+
+        self.assertNotEqual(BOOL, INT)
+        self.assertEqual(REAL, REAL)
+
+        AAIBR = ArrayType(ArrayType(INT, BOOL), REAL)
+        self.assertEqual(str(AAIBR), "Array{Array{Int, Bool}, Real}")
+
+        bt1 = BVType(4)
+        bt2 = BVType(4)
+        self.assertEqual(bt1, bt2)
+        self.assertEqual(bt1.width, 4)
+
+    def test_fake_arrays(self):
+        FakeArrayType = Type("FakeArray", 2)
+        with self.assertRaises(PysmtValueError):
+            FreshSymbol(FakeArrayType)
+        FakeArrayII = FakeArrayType(INT, INT)
+        self.assertEqual(str(FakeArrayII), "FakeArray{Int, Int}")
+        self.assertEqual(FakeArrayII.as_smtlib(False), "(FakeArray Int Int)")
+        s = FreshSymbol(FakeArrayII)
+        self.assertIsNotNone(s)
+
+    def test_simple_sorts(self):
+        # (define-sort I () Int)
+        # (define-sort Set (T) (Array T Bool))
+        I = INT
+        SET = PartialType("Set", lambda t1: ArrayType(t1, BOOL))
+        self.assertEqual(ArrayType(INT, BOOL), SET(I))
+
+        # (declare-const s1 (Set I))
+        # (declare-const a I)
+        # (declare-const b Int)
+        s1 = FreshSymbol(SET(I))
+        a = FreshSymbol(I)
+        b = FreshSymbol(INT)
+
+        # (= (select s1 a) true)
+        # (= (select s1 b) false)
+        f1 = EqualsOrIff(Select(s1, a), TRUE())
+        f2 = EqualsOrIff(Select(s1, b), FALSE())
+        self.assertIsNotNone(f1)
+        self.assertIsNotNone(f2)
+
+        # Cannot instantiate a PartialType directly:
+        with self.assertRaises(PysmtValueError):
+            FreshSymbol(SET)
+
+        # (declare-sort A 0)
+        # Uninterpreted sort
+        A = Type("A", 0)
+        B = Type("B")
+
+        c1 = FreshSymbol(A)
+        c2 = FreshSymbol(A)
+        c3 = FreshSymbol(Type("A"))
+        c4 = FreshSymbol(B)
+        EqualsOrIff(c1, c2)
+        EqualsOrIff(c2, c3)
+        with self.assertRaises(PysmtTypeError):
+            EqualsOrIff(c1, c4)
+
+        with self.assertRaises(PysmtValueError):
+            Type("A", 1)
+
+        C = Type("C", 1)
+        c5 = FreshSymbol(C(A))
+        c6 = FreshSymbol(C(B))
+        self.assertIsNotNone(c5)
+        with self.assertRaises(PysmtTypeError):
+            EqualsOrIff(c5, c6)
+
+        # Nesting
+        ty = C(C(C(C(C(A)))))
+        self.assertIsNotNone(FreshSymbol(ty))
+
+        pty = PartialType("pty", lambda S,T: S(S(S(S(S(T))))))
+        self.assertEqual(pty(C,A), ty)
+
+    def test_normalization(self):
+        from pysmt.environment import Environment
+
+        env2 = Environment()
+        mgr2 = env2.formula_manager
+
+        ty = ArrayType(BOOL, REAL)
+        x = FreshSymbol(ty)
+        fty = FunctionType(BOOL, (ty,))
+        f = FreshSymbol(fty)
+        g = Function(f, (x,))
+        self.assertIsNotNone(g)
+        self.assertNotIn(g, mgr2)
+        g2 = mgr2.normalize(g)
+        self.assertIn(g2, mgr2)
+        # Since the types are from two different environments, they
+        # should be different.
+        x2 = g2.arg(0)
+        ty2 = x2.symbol_type()
+        self.assertFalse(ty2 is ty, ty)
+        fname = g2.function_name()
+        fty2 = fname.symbol_type()
+        self.assertFalse(fty2 is fty, fty)
+
+        # Test ArrayValue
+        h = Array(BVType(4), BV(0, 4))
+        h2 = mgr2.normalize(h)
+        self.assertEqual(h.array_value_index_type(),
+                         h2.array_value_index_type())
+        self.assertFalse(h.array_value_index_type() is \
+                         h2.array_value_index_type())
+
+
+if __name__ == '__main__':
+    main()