src/openpower/visitor.py

   1 #  ___________________________________________________________________________
   2 #
   3 #  Pyomo: Python Optimization Modeling Objects
   4 #  Copyright (c) 2008-2022
   5 #  National Technology and Engineering Solutions of Sandia, LLC
   6 #  Under the terms of Contract DE-NA0003525 with National Technology and
   7 #  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
   8 #  rights in this software.
   9 #  This software is distributed under the 3-clause BSD License.
  10 #  ___________________________________________________________________________
  11
  12 from __future__ import division
  13
  14 import inspect
  15 import logging
  16 import sys
  17 from copy import deepcopy
  18 from collections import deque
  19
  20 logger = logging.getLogger('pyomo.core')
  21
  22 from openpower.numeric_types import nonpyomo_leaf_types
  23
  24
  25 try:
  26     # sys._getframe is slightly faster than inspect's currentframe, but
  27     # is not guaranteed to exist everywhere
  28     currentframe = sys._getframe
  29 except AttributeError:
  30     currentframe = inspect.currentframe
  31
  32
  33 def get_stack_depth():
  34     n = -1  # skip *this* frame in the count
  35     f = currentframe()
  36     while f is not None:
  37         n += 1
  38         f = f.f_back
  39     return n
  40
  41
  42 # For efficiency, we want to run recursively, but don't want to hit
  43 # Python's recursion limit (because that would be difficult to recover
  44 # from cleanly).  However, there is a non-trivial cost to determine the
  45 # current stack depth - and we don't want to hit that for every call.
  46 # Instead, we will assume that the walker is always called with at
  47 # least RECURSION_LIMIT frames available on the stack.  When we hit the
  48 # end of that limit, we will actually check how much space is left on
  49 # the stack and run recursively until only 2*RECURSION_LIMIT frames are
  50 # left.  For the vast majority of well-formed expressions this approach
  51 # avoids a somewhat costly call to get_stack_depth, but still catches
  52 # the vast majority of cases that could generate a recursion error.
  53 RECURSION_LIMIT = 50
  54
  55
  56 class RevertToNonrecursive(Exception):
  57     pass
  58
  59
  60 # NOTE: This module also has dependencies on numeric_expr; however, to
  61 # avoid circular dependencies, we will NOT import them here.  Instead,
  62 # until we can resolve the circular dependencies, they will be injected
  63 # into this module by the .current module (which must be imported
  64 # *after* numeric_expr, logocal_expr, and this module.
  65
  66
  67 # -------------------------------------------------------
  68 #
  69 # Visitor Logic
  70 #
  71 # -------------------------------------------------------
  72
  73
  74 class StreamBasedExpressionVisitor(object):
  75     """This class implements a generic stream-based expression walker.
  76
  77      This visitor walks an expression tree using a depth-first strategy
  78      and generates a full event stream similar to other tree visitors
  79      (e.g., the expat XML parser).  The following events are triggered
  80      through callback functions as the traversal enters and leaves nodes
  81      in the tree:
  82
  83        initializeWalker(expr) -> walk, result
  84        enterNode(N1) -> args, data
  85        {for N2 in args:}
  86          beforeChild(N1, N2) -> descend, child_result
  87            enterNode(N2) -> N2_args, N2_data
  88            [...]
  89            exitNode(N2, n2_data) -> child_result
  90          acceptChildResult(N1, data, child_result) -> data
  91          afterChild(N1, N2) -> None
  92        exitNode(N1, data) -> N1_result
  93        finalizeWalker(result) -> result
  94
  95      Individual event callbacks match the following signatures:
  96
  97     walk, result = initializeWalker(self, expr):
  98
  99          initializeWalker() is called to set the walker up and perform
 100          any preliminary processing on the root node.  The method returns
 101          a flag indicating if the tree should be walked and a result.  If
 102          `walk` is True, then result is ignored.  If `walk` is False,
 103          then `result` is returned as the final result from the walker,
 104          bypassing all other callbacks (including finalizeResult).
 105
 106     args, data = enterNode(self, node):
 107
 108          enterNode() is called when the walker first enters a node (from
 109          above), and is passed the node being entered.  It is expected to
 110          return a tuple of child `args` (as either a tuple or list) and a
 111          user-specified data structure for collecting results.  If None
 112          is returned for args, the node's args attribute is used for
 113          expression types and the empty tuple for leaf nodes.  Returning
 114          None is equivalent to returning (None,None).  If the callback is
 115          not defined, the default behavior is equivalent to returning
 116          (None, []).
 117
 118      node_result = exitNode(self, node, data):
 119
 120          exitNode() is called after the node is completely processed (as
 121          the walker returns up the tree to the parent node).  It is
 122          passed the node and the results data structure (defined by
 123          enterNode() and possibly further modified by
 124          acceptChildResult()), and is expected to return the "result" for
 125          this node.  If not specified, the default action is to return
 126          the data object from enterNode().
 127
 128      descend, child_result = beforeChild(self, node, child, child_idx):
 129
 130          beforeChild() is called by a node for every child before
 131          entering the child node.  The node, child node, and child index
 132          (position in the args list from enterNode()) are passed as
 133          arguments.  beforeChild should return a tuple (descend,
 134          child_result).  If descend is False, the child node will not be
 135          entered and the value returned to child_result will be passed to
 136          the node's acceptChildResult callback.  Returning None is
 137          equivalent to (True, None).  The default behavior if not
 138          specified is equivalent to (True, None).
 139
 140      data = acceptChildResult(self, node, data, child_result, child_idx):
 141
 142          acceptChildResult() is called for each child result being
 143          returned to a node.  This callback is responsible for recording
 144          the result for later processing or passing up the tree.  It is
 145          passed the node, result data structure (see enterNode()), child
 146          result, and the child index (position in args from enterNode()).
 147          The data structure (possibly modified or replaced) must be
 148          returned.  If acceptChildResult is not specified, it does
 149          nothing if data is None, otherwise it calls data.append(result).
 150
 151      afterChild(self, node, child, child_idx):
 152
 153          afterChild() is called by a node for every child node
 154          immediately after processing the node is complete before control
 155          moves to the next child or up to the parent node.  The node,
 156          child node, an child index (position in args from enterNode())
 157          are passed, and nothing is returned.  If afterChild is not
 158          specified, no action takes place.
 159
 160      finalizeResult(self, result):
 161
 162          finalizeResult() is called once after the entire expression tree
 163          has been walked.  It is passed the result returned by the root
 164          node exitNode() callback.  If finalizeResult is not specified,
 165          the walker returns the result obtained from the exitNode
 166          callback on the root node.
 167
 168      Clients interact with this class by either deriving from it and
 169      implementing the necessary callbacks (see above), assigning callable
 170      functions to an instance of this class, or passing the callback
 171      functions as arguments to this class' constructor.
 172
 173     """
 174
 175     # The list of event methods that can either be implemented by
 176     # derived classes or specified as callback functions to the class
 177     # constructor.
 178     #
 179     # This is a dict mapping the callback name to a single character
 180     # that we can use to classify the set of callbacks used by a
 181     # particular Visitor (we define special-purpose node processors for
 182     # certain common combinations).  For example, a 'bex' visitor is one
 183     # that supports beforeChild, enterNode, and exitNode, but NOT
 184     # afterChild or acceptChildResult.
 185     client_methods = {
 186         'enterNode': 'e',
 187         'exitNode': 'x',
 188         'beforeChild': 'b',
 189         'afterChild': 'a',
 190         'acceptChildResult': 'c',
 191         'initializeWalker': '',
 192         'finalizeResult': '',
 193     }
 194
 195     def __init__(self, **kwds):
 196         # This is slightly tricky: We want derived classes to be able to
 197         # override the "None" defaults here, and for keyword arguments
 198         # to override both.  The hasattr check prevents the "None"
 199         # defaults from overriding attributes or methods defined on
 200         # derived classes.
 201         for field in self.client_methods:
 202             if field in kwds:
 203                 setattr(self, field, kwds.pop(field))
 204             elif not hasattr(self, field):
 205                 setattr(self, field, None)
 206         if kwds:
 207             raise RuntimeError("Unrecognized keyword arguments: %s" % (kwds,))
 208
 209         # Handle deprecated APIs
 210         _fcns = (('beforeChild', 2), ('acceptChildResult', 3), ('afterChild', 2))
 211         for name, nargs in _fcns:
 212             fcn = getattr(self, name)
 213             if fcn is None:
 214                 continue
 215             _args = inspect.getfullargspec(fcn)
 216             _self_arg = 1 if inspect.ismethod(fcn) else 0
 217             if len(_args.args) == nargs + _self_arg and _args.varargs is None:
 218                 print(
 219                     "Note that the API for the StreamBasedExpressionVisitor " \
 220                     "has changed to include the child index for the %s() " \
 221                     "method" % (name,))#, file=sys.stderr)
 222
 223                 def wrap(fcn, nargs):
 224                     def wrapper(*args):
 225                         return fcn(*args[:nargs])
 226
 227                     return wrapper
 228
 229                 setattr(self, name, wrap(fcn, nargs))
 230
 231         self.recursion_stack = None
 232
 233         # Set up the custom recursive node handler function (customized
 234         # for the specific set of callbacks that are defined for this
 235         # class instance).
 236         recursive_node_handler = '_process_node_' + ''.join(
 237             sorted(
 238                 '' if getattr(self, f[0]) is None else f[1]
 239                 for f in self.client_methods.items()
 240             )
 241         )
 242         self._process_node = getattr(
 243             self, recursive_node_handler, self._process_node_general
 244         )
 245
 246     def walk_expression(self, expr):
 247         """Walk an expression, calling registered callbacks."""
 248         if self.initializeWalker is not None:
 249             walk, root = self.initializeWalker(expr)
 250             if not walk:
 251                 return root
 252             elif root is None:
 253                 root = expr
 254         else:
 255             root = expr
 256
 257         try:
 258             result = self._process_node(root, RECURSION_LIMIT)
 259             _nonrecursive = None
 260         except RevertToNonrecursive:
 261             ptr = (None,) + self.recursion_stack.pop()
 262             while self.recursion_stack:
 263                 ptr = (ptr,) + self.recursion_stack.pop()
 264             self.recursion_stack = None
 265             _nonrecursive = self._nonrecursive_walker_loop, ptr
 266         except RecursionError:
 267             logger.warning(
 268                 'Unexpected RecursionError walking an expression tree.',
 269                 extra={'id': 'W1003'},
 270             )
 271             _nonrecursive = self.walk_expression_nonrecursive, expr
 272
 273         if _nonrecursive is not None:
 274             return _nonrecursive[0](_nonrecursive[1])
 275
 276         if self.finalizeResult is not None:
 277             return self.finalizeResult(result)
 278         else:
 279             return result
 280
 281     def _compute_actual_recursion_limit(self):
 282         recursion_limit = (
 283             sys.getrecursionlimit() - get_stack_depth() - 2 * RECURSION_LIMIT
 284         )
 285         if recursion_limit <= RECURSION_LIMIT:
 286             self.recursion_stack = []
 287             raise RevertToNonrecursive()
 288         return recursion_limit
 289
 290     def _process_node_general(self, node, recursion_limit):
 291         """Recursive routine for processing nodes with general callbacks
 292
 293         This is the "general" implementation of the
 294         StreamBasedExpressionVisitor node processor that can handle any
 295         combination of registered callback functions.
 296
 297         """
 298         if not recursion_limit:
 299             recursion_limit = self._compute_actual_recursion_limit()
 300         else:
 301             recursion_limit -= 1
 302
 303         if self.enterNode is not None:
 304             tmp = self.enterNode(node)
 305             if tmp is None:
 306                 args = data = None
 307             else:
 308                 args, data = tmp
 309         else:
 310             args = None
 311             data = []
 312         if args is None:
 313             if type(node) in nonpyomo_leaf_types or not node.is_expression_type():
 314                 args = ()
 315             else:
 316                 args = node.args
 317
 318         # Because we do not require the args to be a context manager, we
 319         # will mock up the "with args" using a try-finally.
 320         context_manager = hasattr(args, '__enter__')
 321         if context_manager:
 322             args.__enter__()
 323
 324         try:
 325             descend = True
 326             child_idx = -1
 327             # Note: this relies on iter(iterator) returning the
 328             # iterator.  This seems to hold for all common iterators
 329             # (list, tuple, generator, etc)
 330             arg_iter = iter(args)
 331             for child in arg_iter:
 332                 child_idx += 1
 333                 if self.beforeChild is not None:
 334                     tmp = self.beforeChild(node, child, child_idx)
 335                     if tmp is None:
 336                         descend = True
 337                     else:
 338                         descend, child_result = tmp
 339
 340                 if descend:
 341                     child_result = self._process_node(child, recursion_limit)
 342
 343                 if self.acceptChildResult is not None:
 344                     data = self.acceptChildResult(node, data, child_result, child_idx)
 345                 elif data is not None:
 346                     data.append(child_result)
 347
 348                 if self.afterChild is not None:
 349                     self.afterChild(node, child, child_idx)
 350         except RevertToNonrecursive:
 351             self._recursive_frame_to_nonrecursive_stack(locals())
 352             context_manager = False
 353             raise
 354         finally:
 355             if context_manager:
 356                 args.__exit__(None, None, None)
 357
 358         # We are done with this node.  Call exitNode to compute
 359         # any result
 360         if self.exitNode is not None:
 361             return self.exitNode(node, data)
 362         else:
 363             return data
 364
 365     def _process_node_bex(self, node, recursion_limit):
 366         """Recursive routine for processing nodes with only 'bex' callbacks
 367
 368         This is a special-case implementation of the "general"
 369         StreamBasedExpressionVisitor node processor for the case that
 370         only beforeChild, enterNode, and exitNode are defined (see
 371         also the definition of the client_methods dict).
 372
 373         """
 374         if not recursion_limit:
 375             recursion_limit = self._compute_actual_recursion_limit()
 376         else:
 377             recursion_limit -= 1
 378
 379         tmp = self.enterNode(node)
 380         if tmp is None:
 381             args = data = None
 382         else:
 383             args, data = tmp
 384         if args is None:
 385             if type(node) in nonpyomo_leaf_types or not node.is_expression_type():
 386                 args = ()
 387             else:
 388                 args = node.args
 389
 390         # Because we do not require the args to be a context manager, we
 391         # will mock up the "with args" using a try-finally.
 392         context_manager = hasattr(args, '__enter__')
 393         if context_manager:
 394             args.__enter__()
 395
 396         try:
 397             child_idx = -1
 398             # Note: this relies on iter(iterator) returning the
 399             # iterator.  This seems to hold for all common iterators
 400             # (list, tuple, generator, etc)
 401             arg_iter = iter(args)
 402             for child in arg_iter:
 403                 child_idx += 1
 404                 tmp = self.beforeChild(node, child, child_idx)
 405                 if tmp is None:
 406                     descend = True
 407                 else:
 408                     descend, child_result = tmp
 409
 410                 if descend:
 411                     data.append(self._process_node(child, recursion_limit))
 412                 else:
 413                     data.append(child_result)
 414         except RevertToNonrecursive:
 415             self._recursive_frame_to_nonrecursive_stack(locals())
 416             context_manager = False
 417             raise
 418         finally:
 419             if context_manager:
 420                 args.__exit__(None, None, None)
 421
 422         # We are done with this node.  Call exitNode to compute
 423         # any result
 424         return self.exitNode(node, data)
 425
 426     def _process_node_bx(self, node, recursion_limit):
 427         """Recursive routine for processing nodes with only 'bx' callbacks
 428
 429         This is a special-case implementation of the "general"
 430         StreamBasedExpressionVisitor node processor for the case that
 431         only beforeChild and exitNode are defined (see also the
 432         definition of the client_methods dict).
 433
 434         """
 435         if not recursion_limit:
 436             recursion_limit = self._compute_actual_recursion_limit()
 437         else:
 438             recursion_limit -= 1
 439
 440         if type(node) in nonpyomo_leaf_types or not node.is_expression_type():
 441             args = ()
 442         else:
 443             args = node.args
 444         data = []
 445
 446         try:
 447             child_idx = -1
 448             # Note: this relies on iter(iterator) returning the
 449             # iterator.  This seems to hold for all common iterators
 450             # (list, tuple, generator, etc)
 451             arg_iter = iter(args)
 452             for child in arg_iter:
 453                 child_idx += 1
 454                 tmp = self.beforeChild(node, child, child_idx)
 455                 if tmp is None:
 456                     descend = True
 457                 else:
 458                     descend, child_result = tmp
 459                 if descend:
 460                     data.append(self._process_node(child, recursion_limit))
 461                 else:
 462                     data.append(child_result)
 463         except RevertToNonrecursive:
 464             self._recursive_frame_to_nonrecursive_stack(locals())
 465             raise
 466         finally:
 467             pass
 468
 469         # We are done with this node.  Call exitNode to compute
 470         # any result
 471         return self.exitNode(node, data)
 472
 473     def _recursive_frame_to_nonrecursive_stack(self, local):
 474         child_idx = local['child_idx']
 475         _arg_list = [None] * child_idx
 476         _arg_list.append(local['child'])
 477         _arg_list.extend(local['arg_iter'])
 478         if not self.recursion_stack:
 479             # For the deepest stack frame, the recursion limit hit
 480             # as we started to enter the child.  As we haven't
 481             # started processing it yet, we need to decrement
 482             # child_idx so that it is revisited
 483             child_idx -= 1
 484         self.recursion_stack.append(
 485             (local['node'], _arg_list, len(_arg_list) - 1, local['data'], child_idx)
 486         )
 487
 488     def walk_expression_nonrecursive(self, expr):
 489         """Walk an expression, calling registered callbacks."""
 490         #
 491         # This walker uses a linked list to store the stack (instead of
 492         # an array).  The nodes of the linked list are 6-member tuples:
 493         #
 494         #    ( pointer to parent,
 495         #      expression node,
 496         #      tuple/list of child nodes (arguments),
 497         #      number of child nodes (arguments),
 498         #      data object to aggregate results from child nodes,
 499         #      current child node index )
 500         #
 501         # The walker only needs a single pointer to the end of the list
 502         # (ptr).  The beginning of the list is indicated by a None
 503         # parent pointer.
 504         #
 505         if self.initializeWalker is not None:
 506             walk, result = self.initializeWalker(expr)
 507             if not walk:
 508                 return result
 509             elif result is not None:
 510                 expr = result
 511         if self.enterNode is not None:
 512             tmp = self.enterNode(expr)
 513             if tmp is None:
 514                 args = data = None
 515             else:
 516                 args, data = tmp
 517         else:
 518             args = None
 519             data = []
 520         if args is None:
 521             if type(expr) in nonpyomo_leaf_types or not expr.is_expression_type():
 522                 args = ()
 523             else:
 524                 args = expr.args
 525         if hasattr(args, '__enter__'):
 526             args.__enter__()
 527         node = expr
 528         # Note that because we increment child_idx just before fetching
 529         # the child node, it must be initialized to -1, and ptr[3] must
 530         # always be *one less than* the number of arguments
 531         return self._nonrecursive_walker_loop(
 532             (None, node, args, len(args) - 1, data, -1)
 533         )
 534
 535     def _nonrecursive_walker_loop(self, ptr):
 536         _, node, args, _, data, child_idx = ptr
 537         try:
 538             while 1:
 539                 if child_idx < ptr[3]:
 540                     # Increment the child index pointer here for
 541                     # consistency.  Note that this means that for the bulk
 542                     # of the time, 'child_idx' will not match the value of
 543                     # ptr[5].  This provides a modest performance
 544                     # improvement, as we only have to recreate the ptr tuple
 545                     # just before we descend further into the tree (i.e., we
 546                     # avoid recreating the tuples for the special case where
 547                     # beforeChild indicates that we should not descend
 548                     # further).
 549                     child_idx += 1
 550                     # This node still has children to process
 551                     child = ptr[2][child_idx]
 552
 553                     # Notify this node that we are about to descend into a
 554                     # child.
 555                     if self.beforeChild is not None:
 556                         tmp = self.beforeChild(node, child, child_idx)
 557                         if tmp is None:
 558                             descend = True
 559                             child_result = None
 560                         else:
 561                             descend, child_result = tmp
 562                         if not descend:
 563                             # We are aborting processing of this child node.
 564                             # Tell this node to accept the child result and
 565                             # we will move along
 566                             if self.acceptChildResult is not None:
 567                                 data = self.acceptChildResult(
 568                                     node, data, child_result, child_idx
 569                                 )
 570                             elif data is not None:
 571                                 data.append(child_result)
 572                             # And let the node know that we are done with a
 573                             # child node
 574                             if self.afterChild is not None:
 575                                 self.afterChild(node, child, child_idx)
 576                             # Jump to the top to continue processing the
 577                             # next child node
 578                             continue
 579
 580                     # Update the child argument counter in the stack.
 581                     # Because we are using tuples, we need to recreate the
 582                     # "ptr" object (linked list node)
 583                     ptr = ptr[:4] + (data, child_idx)
 584
 585                     # We are now going to actually enter this node.  The
 586                     # node will tell us the list of its child nodes that we
 587                     # need to process
 588                     if self.enterNode is not None:
 589                         tmp = self.enterNode(child)
 590                         if tmp is None:
 591                             args = data = None
 592                         else:
 593                             args, data = tmp
 594                     else:
 595                         args = None
 596                         data = []
 597                     if args is None:
 598                         if (
 599                             type(child) in nonpyomo_leaf_types
 600                             or not child.is_expression_type()
 601                         ):
 602                             # Leaves (either non-pyomo types or
 603                             # non-Expressions) have no child arguments, so
 604                             # are just put on the stack
 605                             args = ()
 606                         else:
 607                             args = child.args
 608                     if hasattr(args, '__enter__'):
 609                         args.__enter__()
 610                     node = child
 611                     child_idx = -1
 612                     ptr = (ptr, node, args, len(args) - 1, data, child_idx)
 613
 614                 else:  # child_idx == ptr[3]:
 615                     # We are done with this node.  Call exitNode to compute
 616                     # any result
 617                     if hasattr(ptr[2], '__exit__'):
 618                         ptr[2].__exit__(None, None, None)
 619                     if self.exitNode is not None:
 620                         node_result = self.exitNode(node, data)
 621                     else:
 622                         node_result = data
 623
 624                     # Pop the node off the linked list
 625                     ptr = ptr[0]
 626                     # If we have returned to the beginning, return the final
 627                     # answer
 628                     if ptr is None:
 629                         if self.finalizeResult is not None:
 630                             return self.finalizeResult(node_result)
 631                         else:
 632                             return node_result
 633                     # Not done yet, update node to point to the new active
 634                     # node
 635                     node, child = ptr[1], node
 636                     data = ptr[4]
 637                     child_idx = ptr[5]
 638
 639                     # We need to alert the node to accept the child's result:
 640                     if self.acceptChildResult is not None:
 641                         data = self.acceptChildResult(
 642                             node, data, node_result, child_idx
 643                         )
 644                     elif data is not None:
 645                         data.append(node_result)
 646
 647                     # And let the node know that we are done with a child node
 648                     if self.afterChild is not None:
 649                         self.afterChild(node, child, child_idx)
 650
 651         finally:
 652             while ptr is not None:
 653                 if hasattr(ptr[2], '__exit__'):
 654                     ptr[2].__exit__(None, None, None)
 655                 ptr = ptr[0]
 656
 657